Custom op to update cache for torch.cond

[larryliu0820]

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:

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:

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),
    )

[larryliu0820]

Stack from ghstack (oldest at bottom):

• Move torch.cond predicate non-persistent buffer to CPU #16378
• -> Custom op to update cache for torch.cond #16366

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/16366

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 3 New Failures, 26 Pending, 1 Unrelated Failure

As of commit a8b20f5 with merge base 1ce615e ():

NEW FAILURES - The following jobs have failed:

• pull / test-samsung-models-linux / linux-job (gh)
  RuntimeError: Command docker exec -t fee8e70fe8a6d9fedb0bbdefe546ba49994968959e476fb0c43895f1c832e0a0 /exec failed with exit code 1
• pull / test-samsung-quantmodels-linux / linux-job (gh)
  RuntimeError: Command docker exec -t 8bc84b1073531f98c1af29bb04de023427d283c49b6b2582f684f88ac6cd8f37 /exec failed with exit code 1
• trunk / test-mcu-cortex-m-backend / linux-job (gh)
  RuntimeError: Command docker exec -t 778c8f03990336156be0a84f2ea769564b72931c940fbd1e5239b6059d995e3c /exec failed with exit code 1

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

• pull / android / run-emulator (gh) (#16137)
  Timeout waiting for emulator to boot.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[mergennachin]

This is returning the original cache whereas eager mode returns a copy. Do you want this discrepancy?

[larryliu0820]

yeah. eager has to return a copy to make exported_program runnable. For inductor lowering we need it to be as efficient as possible.

[mergennachin]

Okay then document the discrepancy