Guard the sliding-window ring KV cache: concat on prefill, write-first on decode (#314)

python/src/coreai_models/models/macos/muse_glimmer.py

@@ -18,6 +18,10 @@
 KV cache: split into sliding (ring buffer, fixed size) and global (growing).
 Sliding layers use RingKVCache at window_size capacity; global layers use
 standard KVCache with dynamic sequence length.
+
+Sliding attention is guarded by ``prefill_mode``: prefill reads before writing
+(concat, correct for any query_len); decode writes first (fast, query_len == 1
+only -- the Swift engine rejects query_len > 1 on this path).
 """
 
 import gc
@@ -44,13 +48,38 @@
     _load_tensors_for_keys,
     _resolve_safetensors_files,
 )
-from coreai_models.primitives.macos.cache import KVCache, RingKVCache, ring_window_causal_mask
+from coreai_models.primitives.macos.cache import (
+    KVCache,
+    RingKVCache,
+    concat_window_causal_mask,
+    ring_window_causal_mask,
+)
 from coreai_models.primitives.macos.mlp import MLP
 from coreai_models.primitives.macos.rms_norm import RMSNorm, RMSNormPlusOne
 from coreai_models.primitives.macos.rope import RoPE
 from coreai_models.primitives.macos.sdpa import SDPA
 
 
+def _sliding_window_mask(
+    prefill_mode: bool,
+    query_len: int,
+    capacity: int,
+    offset: int,
+    window_size: int,
+    device: torch.device,
+) -> torch.Tensor:
+    """Pick the sliding-window mask that matches the ring access pattern.
+
+    ``prefill_mode`` picks the matching pair: prefill uses fetch_and_concat +
+    concat_window_causal_mask (correct after the ring wraps); decode uses
+    update_and_fetch + ring_window_causal_mask (cheaper, query_len == 1 only,
+    Swift guards query_len > 1). Resolved at trace time.
+    """
+    if prefill_mode:
+        return concat_window_causal_mask(query_len, capacity, offset, window_size, device)
+    return ring_window_causal_mask(query_len, capacity, offset, window_size, device)
+
+
 class Attention(nn.Module):
     def __init__(self, config, layer_idx: int) -> None:
         super().__init__()
@@ -99,6 +128,7 @@ def forward(
         global_cache: KVCache | None = None,
         sliding_kv_cache: RingKVCache | None = None,
         sliding_mask: torch.Tensor | None = None,
+        prefill_mode: bool = False,
     ) -> torch.Tensor:
         batch_size = x.shape[0]
         n_heads, n_kv_heads = self.n_heads, self.n_kv_heads
@@ -132,10 +162,27 @@ def forward(
             key = self.rope(key, position_ids=rope_positions, freqs=freqs)
 
         if self.is_sliding and sliding_kv_cache is not None:
-            key, value = sliding_kv_cache.update_and_fetch(
-                self.cache_slot_idx, offset, key, value, query_len=query_len
-            )
-            attn_output = self.sdpa(query=query, key=key, value=value, attn_mask=sliding_mask)
+            if prefill_mode:
+                # Prefill: read before write (concat) -- writing first would
+                # corrupt a multi-token chunk that still needs the evicted history.
+                full_key, full_value = sliding_kv_cache.fetch_and_concat(
+                    self.cache_slot_idx, key, value
+                )
+                attn_output = self.sdpa(
+                    query=query, key=full_key, value=full_value, attn_mask=sliding_mask
+                )
+                sliding_kv_cache.update(
+                    self.cache_slot_idx, offset, key, value, query_len=query_len
+                )
+            else:
+                # Decode: write first (cheaper). Correct for query_len == 1 only;
+                # Swift rejects query_len > 1 on this path.
+                full_key, full_value = sliding_kv_cache.update_and_fetch(
+                    self.cache_slot_idx, offset, key, value, query_len=query_len
+                )
+                attn_output = self.sdpa(
+                    query=query, key=full_key, value=full_value, attn_mask=sliding_mask
+                )
         elif not self.is_sliding and global_cache is not None:
             key, value = global_cache.update_and_fetch(
                 self.cache_slot_idx, offset, key, value, seq_len=seq_len, query_len=query_len
@@ -173,6 +220,7 @@ def forward(
         global_cache: KVCache | None = None,
         sliding_kv_cache: RingKVCache | None = None,
         sliding_mask: torch.Tensor | None = None,
+        prefill_mode: bool = False,
     ) -> torch.Tensor:
         r = self.self_attn(
             self.input_layernorm(x),
@@ -182,6 +230,7 @@ def forward(
             global_cache,
             sliding_kv_cache,
             sliding_mask,
+            prefill_mode,
         )
         r = self.post_attention_layernorm(r)
         h = x + r
@@ -222,6 +271,7 @@ def forward(
         position_ids: torch.IntTensor,
         global_cache: KVCache | None = None,
         sliding_kv_cache: RingKVCache | None = None,
+        prefill_mode: bool = False,
     ) -> torch.Tensor:
         query_len = input_ids.shape[-1]
         seq_len = position_ids.shape[-1]
@@ -235,7 +285,8 @@ def forward(
 
         sliding_mask = None
         if sliding_kv_cache is not None:
-            sliding_mask = ring_window_causal_mask(
+            sliding_mask = _sliding_window_mask(
+                prefill_mode,
                 query_len,
                 sliding_kv_cache.capacity(),
                 offset,
@@ -252,6 +303,7 @@ def forward(
                 global_cache,
                 sliding_kv_cache,
                 sliding_mask,
+                prefill_mode,
             )
         h = self.norm(h)
         h = h * self.output_multiplier
@@ -411,7 +463,9 @@ def forward(
     ) -> torch.Tensor | tuple:
         global_cache = KVCache(global_k_cache, global_v_cache)
         sliding_kv_cache = RingKVCache(sliding_k_cache, sliding_v_cache)
-        out = self.model(input_ids, position_ids, global_cache, sliding_kv_cache)
+        out = self.model(
+            input_ids, position_ids, global_cache, sliding_kv_cache, prefill_mode=self.prefill_mode
+        )
         if self.prefill_mode:
             return ()
         logits = self.lm_head(out)
@@ -577,6 +631,7 @@ def forward(  # type: ignore[override]
         position_ids: torch.IntTensor,
         global_cache: KVCache | None = None,
         sliding_cache: RingKVCache | None = None,
+        prefill_mode: bool = False,
     ) -> tuple[torch.Tensor, list[torch.Tensor]]:
         query_len = input_ids.shape[-1]
         seq_len = position_ids.shape[-1]
@@ -590,7 +645,8 @@ def forward(  # type: ignore[override]
 
         sliding_mask = None
         if sliding_cache is not None:
-            sliding_mask = ring_window_causal_mask(
+            sliding_mask = _sliding_window_mask(
+                prefill_mode,
                 query_len,
                 sliding_cache.capacity(),
                 offset,
@@ -609,6 +665,7 @@ def forward(  # type: ignore[override]
                 global_cache,
                 sliding_cache,
                 sliding_mask,
+                prefill_mode,
             )
             if idx in self._tap_set:
                 extracted.append(h)
@@ -648,7 +705,9 @@ def forward(
     ) -> tuple[torch.Tensor, torch.Tensor]:
         global_cache = KVCache(global_k_cache, global_v_cache)
         sliding_cache = RingKVCache(sliding_k_cache, sliding_v_cache)
-        hidden, extracted = self.model(input_ids, position_ids, global_cache, sliding_cache)
+        hidden, extracted = self.model(
+            input_ids, position_ids, global_cache, sliding_cache, prefill_mode=self.prefill_mode
+        )
 
         logits = self.lm_head(hidden)
         if self._softcap:

python/src/coreai_models/models/macos/muse_glimmer_drafter_dflash.py

@@ -160,7 +160,8 @@ def forward(
         query = self.rope(query, position_ids=rope_positions, freqs=freqs)
         key = self.rope(key, position_ids=rope_positions, freqs=freqs)
 
-        # Write draft KV to cache, fetch full cache (injected + draft)
+        # Write draft KV, then fetch. Write-first is safe here (unlike the target):
+        # the block-level dflash mask isn't per-query. Parity harness, not exported.
         key, value = cache.update_and_fetch(self.layer_idx, offset, key, value, query_len=query_len)
 
         attn_output = (

python/src/coreai_models/models/macos/muse_glimmer_drafter_ring.py

@@ -13,6 +13,9 @@
 - No gated attention, no sandwich norms
 - Shares embed_tokens and lm_head with the target model
 - Ring buffer KV cache with explicit attention mask
+
+Sliding attention reads before writing (concat), so a multi-token chunk can't
+evict window history it still needs.
 """
 
 import gc
@@ -38,7 +41,7 @@
     _load_tensors_for_keys,
     _resolve_safetensors_files,
 )
-from coreai_models.primitives.macos.cache import RingKVCache, ring_window_causal_mask
+from coreai_models.primitives.macos.cache import RingKVCache, concat_window_causal_mask
 from coreai_models.primitives.macos.mlp import MLP
 from coreai_models.primitives.macos.rms_norm import RMSNorm
 from coreai_models.primitives.macos.rope import RoPE
@@ -111,15 +114,20 @@ def forward(
         key = self.rope(key, position_ids=rope_positions, freqs=freqs)
 
         if cache is not None:
-            key, value = cache.update_and_fetch(
-                self.layer_idx, offset, key, value, query_len=query_len
-            )
+            # Read before write (concat), else a multi-token chunk garbles output.
+            attn_key, attn_value = cache.fetch_and_concat(self.layer_idx, key, value)
+        else:
+            attn_key, attn_value = key, value
 
         attn_output = (
-            self.sdpa(query=query, key=key, value=value, attn_mask=attn_mask)
+            self.sdpa(query=query, key=attn_key, value=attn_value, attn_mask=attn_mask)
             .permute(0, 2, 1, 3)
             .reshape(batch_size, query_len, self.n_heads * self.head_dim)
         )
+
+        if cache is not None:
+            cache.update(self.layer_idx, offset, key, value, query_len=query_len)
+
         return self.o_proj(attn_output)
 
 
@@ -175,7 +183,9 @@ def forward(
         h = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + self.config.rms_norm_eps)
 
         if cache is not None:
-            attn_mask = ring_window_causal_mask(
+            # Drafter is always multi-token (one main graph, no decode split), so
+            # it stays on the read-before-write concat path.
+            attn_mask = concat_window_causal_mask(
                 query_len=query_len,
                 capacity=cache.capacity(),
                 offset=offset,

python/src/coreai_models/primitives/macos/cache.py

@@ -295,7 +295,10 @@ def ring_window_causal_mask(
     """Sliding-window causal mask for a ring-buffer KV cache.
 
     Reconstructs each slot's absolute position from the ring layout and applies
-    causal + window + validity predicates. Companion to RingKVCache.
+    causal + window + validity predicates. Companion to
+    :meth:`RingKVCache.update_and_fetch`, i.e. the new K/V is assumed to be
+    already written into the ring. Only correct for ``query_len == 1``; for
+    multi-token steps use :func:`concat_window_causal_mask` instead.
 
     Args:
         query_len: number of query positions in this forward call
@@ -328,12 +331,50 @@ def ring_window_causal_mask(
     return causal & in_window & nonneg
 
 
+def concat_window_causal_mask(
+    query_len: int,
+    capacity: int,
+    offset: int,
+    window_size: int,
+    device: torch.device,
+) -> torch.Tensor:
+    """Sliding-window causal mask for ``[cached ring ++ new keys]``, before the write.
+
+    Companion to :meth:`RingKVCache.fetch_and_concat`: ``capacity`` ring slots at
+    positions ``< offset``, then ``query_len`` new keys at ``offset..``. Reading
+    before writing is what keeps a multi-token step correct.
+
+    Returns a bool mask ``[query_len, capacity + query_len]`` (True = attend);
+    columns ``[0, capacity)`` are ring slots, the rest are the new keys.
+    """
+    row = torch.arange(query_len, device=device)
+    q_pos = row.unsqueeze(-1) + offset  # (query_len, 1)
+
+    # Reconstruct each ring slot's absolute position by walking back from offset-1
+    # (scalar modulo only, no tensor remainder; +capacity keeps it non-negative at offset 0).
+    slot = torch.arange(capacity, device=device)  # (capacity,)
+    last_pos = offset - 1
+    r = (offset + capacity - 1) % capacity  # == last_pos % capacity
+    diff = r - slot  # (capacity,), range (-capacity, capacity)
+    ring_back = torch.where(diff >= 0, diff, diff + capacity)
+    k_pos = (last_pos - ring_back).unsqueeze(0)  # (1, capacity)
+
+    # Cached positions are all < offset, so causality is implicit; drop unwritten slots.
+    ring_mask = ((q_pos - k_pos) < window_size) & (k_pos >= 0)
+
+    new_pos = (row + offset).unsqueeze(0)  # (1, query_len)
+    new_mask = (new_pos <= q_pos) & ((q_pos - new_pos) < window_size)
+
+    return torch.cat([ring_mask, new_mask], dim=-1)
+
+
 class RingKVCache:
-    """Ring-buffer KV cache for sliding window attention layers.
+    """Fixed-size ring KV cache for sliding-window layers; writes at position % capacity.
 
-    Fixed-size cache [n_layers, 1, n_kv_heads, capacity, head_dim] that writes
-    at position % capacity and never grows. Use ring_window_causal_mask() to
-    build the attention mask.
+    Two patterns: ``fetch_and_concat`` -> attend -> ``update`` (read before write,
+    any query_len, uses :func:`concat_window_causal_mask`); or ``update_and_fetch``
+    -> attend (write first, cheaper, query_len == 1 only, uses
+    :func:`ring_window_causal_mask`).
     """
 
     def __init__(self, k_cache: torch.Tensor, v_cache: torch.Tensor) -> None:
@@ -343,24 +384,38 @@ def __init__(self, k_cache: torch.Tensor, v_cache: torch.Tensor) -> None:
     def capacity(self) -> int:
         return self._k_cache.shape[3]
 
-    def update_and_fetch(
+    def fetch_and_concat(
         self,
         layer_idx: int,
-        offset: int,
         k: torch.Tensor,
         v: torch.Tensor,
-        query_len: int | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Write K/V at ring position, return full cache for this layer.
+        """Cached window with ``k``/``v`` appended (does not mutate; call
+        :meth:`update` after attention). Pairs with concat_window_causal_mask.
+        """
+        torch._check_is_size(layer_idx)
+        torch._check(layer_idx < self._k_cache.size(0))
+
+        cached_k = self._k_cache.narrow(0, layer_idx, 1).squeeze(0)
+        cached_v = self._v_cache.narrow(0, layer_idx, 1).squeeze(0)
+        # cat copies, so results don't alias the cache (safe across update()).
+        return torch.cat([cached_k, k], dim=-2), torch.cat([cached_v, v], dim=-2)
 
-        The write must not wrap around the ring boundary, i.e.
-        ``(offset % capacity) + query_len <= capacity`` must hold.
-        Decode (query_len=1) always satisfies this. For chunked prefill,
-        choose chunk sizes that align to the capacity boundary (e.g.
-        capacity // 2) so that no single chunk straddles the wrap point.
+    def update(
+        self,
+        layer_idx: int,
+        offset: int,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        query_len: int | None = None,
+    ) -> None:
+        """Write K/V into the ring at ``offset % capacity``. Must not wrap:
+        ``offset % capacity + query_len <= capacity`` (decode always holds; for
+        chunked prefill pick a chunk size dividing ``capacity``). Can't branch on
+        the bound -- ``offset`` is a symint under torch.export.
 
         Raises:
-            RuntimeError: If the write would wrap around the ring buffer.
+            RuntimeError: if the write would wrap the ring.
         """
         if query_len is None:
             query_len = k.shape[-2]
@@ -401,6 +456,23 @@ def update_and_fetch(
                 ),
             )
 
+    def update_and_fetch(
+        self,
+        layer_idx: int,
+        offset: int,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        query_len: int | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Write K/V at the ring position, then return the full cache for this layer.
+
+        Only correct for ``query_len == 1`` — see the class docstring. Prefer
+        :meth:`fetch_and_concat` + :meth:`update` for multi-token steps.
+
+        Raises:
+            RuntimeError: If the write would wrap around the ring buffer.
+        """
+        self.update(layer_idx, offset, k, v, query_len=query_len)
         k_out = self._k_cache.narrow(0, layer_idx, 1).squeeze(0)
         v_out = self._v_cache.narrow(0, layer_idx, 1).squeeze(0)
         return k_out, v_out