Improve YOCO static attention: reusable helper, correct tensor op, runtime guard (pytorch#18545)

Differential Revision: D97637849

Pull Request resolved: pytorch#18545

Author: viveknayakatmeta

examples/models/llama/static_attention.py

@@ -220,6 +220,18 @@ def unmask(self, new_unmasked_len):
         self.unmasked_len += new_unmasked_len
 
 
+def _is_kv_shared_layer(
+    layer_idx: int, n_layers: int, num_kv_shared_layers: int
+) -> bool:
+    """Check if this layer uses shared K/V from a donor layer (YOCO)."""
+    if num_kv_shared_layers <= 0:
+        return False
+    first_shared = n_layers - num_kv_shared_layers
+    if first_shared <= 0:
+        return False
+    return layer_idx >= first_shared
+
+
 class StaticAttentionIOManager:
     class NGramCache:
         def __init__(self, max_size):
@@ -286,11 +298,7 @@ def _from_config(
         self.freqs_sin = freqs[1].to(dtype)
 
         split_mha = config.attention_type in ("static", "static_shas")
-        # YOCO: skip cache creation for KV-shared layers (2nd half)
         num_kv_shared = getattr(config, "num_kv_shared_layers", 0)
-        first_kv_shared = (
-            config.n_layers - num_kv_shared if num_kv_shared > 0 else config.n_layers
-        )
         if split_mha:
             self.k_caches = {
                 StaticKVCache.calculate_cache_key(layer_id, head_id): torch.zeros(
@@ -301,7 +309,8 @@ def _from_config(
                 )
                 for layer_id in range(config.n_layers)
                 for head_id in range(none_throws(config.n_kv_heads))
-                if cache_lens[layer_id] > 0 and layer_id < first_kv_shared
+                if cache_lens[layer_id] > 0
+                and not _is_kv_shared_layer(layer_id, config.n_layers, num_kv_shared)
             }
             self.v_caches = {
                 StaticKVCache.calculate_cache_key(layer_id, head_id): torch.zeros(
@@ -312,7 +321,8 @@ def _from_config(
                 )
                 for layer_id in range(config.n_layers)
                 for head_id in range(none_throws(config.n_kv_heads))
-                if cache_lens[layer_id] > 0 and layer_id < first_kv_shared
+                if cache_lens[layer_id] > 0
+                and not _is_kv_shared_layer(layer_id, config.n_layers, num_kv_shared)
             }
         else:
             self.k_caches = {
@@ -324,7 +334,8 @@ def _from_config(
                     dtype=dtype,
                 )
                 for layer_id in range(config.n_layers)
-                if cache_lens[layer_id] > 0 and layer_id < first_kv_shared
+                if cache_lens[layer_id] > 0
+                and not _is_kv_shared_layer(layer_id, config.n_layers, num_kv_shared)
             }
             self.v_caches = {
                 StaticKVCache.calculate_cache_key(layer_id, 0): torch.zeros(
@@ -335,7 +346,8 @@ def _from_config(
                     dtype=dtype,
                 )
                 for layer_id in range(config.n_layers)
-                if cache_lens[layer_id] > 0 and layer_id < first_kv_shared
+                if cache_lens[layer_id] > 0
+                and not _is_kv_shared_layer(layer_id, config.n_layers, num_kv_shared)
             }
 
         self.generate_full_logits = config.generate_full_logits
@@ -998,8 +1010,10 @@ def forward(
         new_qs = [wq(x) for wq in self.wqs]
 
         shared_kv = kwargs.get("shared_kv")
-        if self.is_kv_shared_layer:
-            assert shared_kv is not None
+        if shared_kv is not None:
+            assert (
+                self.is_kv_shared_layer
+            ), "shared_kv provided but this is not a KV shared layer"
             new_ks = []
             new_vs = []
         else:
@@ -1222,7 +1236,7 @@ def _process_normal_kv(
 
         kv_to_share = None
         if self.num_kv_shared_layers > 0:
-            kv_to_share = (torch.cat(all_ks, dim=1), torch.cat(all_vs, dim=1))
+            kv_to_share = (torch.stack(all_ks, dim=1), torch.stack(all_vs, dim=1))
 
         return torch.cat(heads, dim=-1), out_cache_state, kv_to_share