Custom op to update cache for torch.cond (#16366)

torch.cond doesn't take aliasing or mutations. Adding 2 ops for
supporting conditionally updating kv cache:

* `executorch::alias`: takes 2 tensors and return the same 2 tensors.
* `executorch::update_cross_attn_cache`: takes a tensor `cache` and a
tensor `value`, in place copy `value` into `cache`.

With these 2 ops, we can rewrite the model definition from:

```py
if is_cross_attention and past_key_values and is_updated:
    # reuse k,v, cross_attentions
    key_states = past_key_values.layers[self.layer_idx].keys
    value_states = past_key_values.layers[self.layer_idx].values
else:
    key_states = self.k_proj(current_states).view(bsz, -1, self.num_heads, self.head_dim)
    value_states = self.v_proj(current_states).view(bsz, -1, self.num_heads, self.head_dim)
    key_states = key_states.transpose(1, 2).contiguous()
    value_states = value_states.transpose(1, 2).contiguous()
    if past_key_values is not None:
        # save all key/value_states to cache to be re-used for fast auto-regressive generation
        cache_position = cache_position if not is_cross_attention else None
        key_states, value_states = past_key_values.update(
            key_states, value_states, self.layer_idx, {"cache_position": cache_position}
        )
```

Into:

```py
def use_cached_kv(
    cached_keys: Tensor,
    cached_values: Tensor,
    key_value_states: Tensor,
) -> tuple[Tensor, Tensor]:
    # Just reuse cached K/V
    return torch.ops.executorch.alias(cached_keys, cached_values)

def recompute_kv(
    cached_keys: Tensor,  # unused
    cached_values: Tensor,  # unused
    key_value_states: Tensor,
) -> tuple[Tensor, Tensor]:
    # Compute fresh K/V (export-friendly: use custom op to mutate cache)
    key_states = self.k_proj(key_value_states).view(bsz, -1, self.num_heads, self.head_dim)
    value_states = self.v_proj(key_value_states).view(bsz, -1, self.num_heads, self.head_dim)
    key_states = key_states.transpose(1, 2).contiguous()
    value_states = value_states.transpose(1, 2).contiguous()
    k = torch.ops.executorch.update_cross_attn_cache(key_states, cached_keys)
    v = torch.ops.executorch.update_cross_attn_cache(value_states, cached_values)
    return k, v

if past_key_values is not None and self.layer_idx is not None:
    # Grab cached tensors (these are Tensors, so they are OK for export)
    cached_keys = past_key_values.layers[self.layer_idx].keys
    cached_values = past_key_values.layers[self.layer_idx].values

    # Tensor predicate: True if any element is non-zero
    # Result is a 0-dim bool tensor suitable for torch.cond
    cache_is_initialized = (cached_keys != 0).any()

    # Use torch.cond to select branch in a traceable way.
    # All operands must be (nested) tensors or simple Python values.
    key_states, value_states = torch.cond(
        cache_is_initialized,
        use_cached_kv,
        recompute_kv,
        operands=(cached_keys, cached_values, key_value_states),
    )
```

Author: larryliu0820

.ci/docker/ci_commit_pins/pytorch.txt

@@ -1 +1 @@
-7a064ed3eafa43f17412d434b395240c727b3000
+7a79b41e29a790ebb4b530eb98a89381e2d7de29

extension/llm/custom_ops/TARGETS

@@ -63,3 +63,16 @@ runtime.python_test(
         "//executorch/extension/pybindings:portable_lib",
     ],
 )
+
+runtime.python_test(
+    name = "test_update_cross_attn_cache",
+    srcs = [
+        "test_update_cross_attn_cache.py",
+    ],
+    preload_deps = [
+        ":custom_ops_aot_py",
+    ],
+    deps = [
+        "//caffe2:torch",
+    ],
+)

extension/llm/custom_ops/custom_ops.py

@@ -12,10 +12,16 @@
 
 import logging
 
+from typing import Tuple
+
 import torch
 
+from torch._inductor.lowering import lowerings as L, register_lowering
+
 from torch.library import impl
 
+aten = torch.ops.aten
+
 try:
     op = torch.ops.llama.sdpa_with_kv_cache.default
     assert op is not None
@@ -387,3 +393,103 @@ def custom_quantized_sdpa_meta(
     )
 
     return torch.empty(query.size(), dtype=torch.float32, device="meta")
+
+
+# 1) Define the custom op in the "executorch" namespace with name "alias"
+@torch.library.custom_op("executorch::alias", mutates_args=())
+def custom_alias(x: torch.Tensor, y: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    # no copies, just pass-through
+    return x, y
+
+
+# 2) FakeTensor kernel: describes output metadata for compile-time
+@custom_alias.register_fake
+def _(x, y):
+    # For this op, outputs have exactly the same shape/dtype/device as inputs.
+    # We just need *dummy* tensors with that metadata.
+    out_x = torch.empty_like(x)
+    out_y = torch.empty_like(y)
+    return out_x, out_y
+
+
+@register_lowering(torch.ops.executorch.alias.default)
+def lowering_custom_alias(x, y):
+    # x, y here are IR values (Inductor's internal representation).
+    # Alias is logically a no-op – just pass them through.
+    return x, y
+
+
+# Expecting cache shape: (B, H, S_max, D), value shape (B, H, S, D) where S <= S_max
+def _validate_cross_attn_cache_params(value: torch.Tensor, cache: torch.Tensor):
+    torch._assert(value.dim() == 4, "value must be 4D")
+    torch._assert(cache.dim() == 4, "cache must be 4D")
+    # Cache shape: (B, H, S_max, D)
+    # Value shape: (B, H, S, D)
+    torch._assert(
+        value.size(2) <= cache.size(2),
+        f"value sequence length {value.size(2)} exceeds cache size {cache.size(2)}",
+    )
+    torch._assert(value.size(0) == cache.size(0), "batch size mismatch")
+    torch._assert(value.size(1) == cache.size(1), "num heads mismatch")
+    torch._assert(value.size(3) == cache.size(3), "head dim mismatch")
+    torch._assert(value.dtype == cache.dtype, "dtype mismatch")
+
+
+# Intentionally declaring no mutations to enable use inside torch.cond branches,
+# which require pure functions. torch.cond requires branch functions to be mutation-free.
+# We omit `cache` from `mutates_args` to satisfy this constraint, accepting the
+# mutation for inference use.
+@torch.library.custom_op("executorch::update_cross_attn_cache", mutates_args=[])
+def _update_cross_attn_cache(value: torch.Tensor, cache: torch.Tensor) -> torch.Tensor:
+    """
+    Update cross-attention KV cache with new values.
+
+    Copies the value tensor into the beginning of the cache tensor along the
+    sequence dimension. This is used for cross-attention caching where the
+    encoder outputs are computed once and reused across decoding steps.
+
+    Args:
+        value: New values to store in cache. Shape: [B, H, S, D] where
+            B = batch size, H = num heads, S = sequence length, D = head dim.
+        cache: Pre-allocated cache tensor to update. Shape: [B, H, S_max, D]
+            where S_max >= S.
+
+    Returns:
+        A clone of the updated cache tensor. Note that this is different from
+        inductor lowering which returns the cache tensor itself. The reason is
+        that if we return input buffer directly, we will fail torch check in
+        higher order ops.
+
+    Note:
+        The cache is mutated in-place, but we return a clone to avoid aliasing
+        issues with the exported program.
+    """
+    _validate_cross_attn_cache_params(value, cache)
+    cache[:, :, : value.size(2), :].copy_(value)
+    return cache.clone()
+
+
+# Register the fake (meta) kernel
+@_update_cross_attn_cache.register_fake
+def _update_cross_attn_cache_fake(
+    value: torch.Tensor, cache: torch.Tensor
+) -> torch.Tensor:
+    _validate_cross_attn_cache_params(value, cache)
+    return torch.empty_like(cache)
+
+
+# Register Inductor lowering
+@register_lowering(torch.ops.executorch.update_cross_attn_cache)
+def _update_cross_attn_cache_lowering(value, cache):
+    # cache shape: [B, H, S_max, D]
+    # value shape: [B, H, S, D]
+
+    # We need to slice the cache along dim 2 (sequence length)
+    # slice(self, dim, start, end, step=1)
+    seq_len = value.get_size()[2]
+    cache_slice = L[aten.slice.Tensor](cache, 2, 0, seq_len, 1)
+
+    # Copy value into the slice
+    L[aten.copy_.default](cache_slice, value)
+
+    return cache