diff --git a/python/src/coreai_models/models/macos/muse_glimmer.py b/python/src/coreai_models/models/macos/muse_glimmer.py index 95c4585b..84eec242 100644 --- a/python/src/coreai_models/models/macos/muse_glimmer.py +++ b/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: diff --git a/python/src/coreai_models/models/macos/muse_glimmer_drafter_dflash.py b/python/src/coreai_models/models/macos/muse_glimmer_drafter_dflash.py index 7645d696..3b331a93 100644 --- a/python/src/coreai_models/models/macos/muse_glimmer_drafter_dflash.py +++ b/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 = ( diff --git a/python/src/coreai_models/models/macos/muse_glimmer_drafter_ring.py b/python/src/coreai_models/models/macos/muse_glimmer_drafter_ring.py index 579f27bd..8b460d76 100644 --- a/python/src/coreai_models/models/macos/muse_glimmer_drafter_ring.py +++ b/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, diff --git a/python/src/coreai_models/primitives/macos/cache.py b/python/src/coreai_models/primitives/macos/cache.py index ce1926d0..848b2049 100644 --- a/python/src/coreai_models/primitives/macos/cache.py +++ b/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 diff --git a/python/tests/test_model_units/test_models/test_macos_layers/test_ring_buffer_parity.py b/python/tests/test_model_units/test_models/test_macos_layers/test_ring_buffer_parity.py index 313ea8d5..252cae20 100644 --- a/python/tests/test_model_units/test_models/test_macos_layers/test_ring_buffer_parity.py +++ b/python/tests/test_model_units/test_models/test_macos_layers/test_ring_buffer_parity.py @@ -8,15 +8,26 @@ Validates that incremental decode with the ring buffer produces the same logits as full recompute (ground truth), both within and across the window boundary (wrap-around). Uses chunked prefill to stay within buffer capacity. + +TestConcatWindowAttention covers the read-before-write path used by +muse_glimmer: fetch_and_concat -> attention -> update, which (unlike +update_and_fetch) is correct for multi-token prefill chunks. """ +import math from types import SimpleNamespace import pytest import torch +from coreai_models.models.macos.muse_glimmer import MuseGlimmerModel from coreai_models.models.macos.muse_glimmer_drafter_ring import DrafterRingModel -from coreai_models.primitives.macos.cache import RingKVCache, ring_window_causal_mask +from coreai_models.primitives.macos.cache import ( + KVCache, + RingKVCache, + concat_window_causal_mask, + ring_window_causal_mask, +) def _drafter_config(window: int = 64) -> SimpleNamespace: @@ -249,3 +260,535 @@ def test_mask_after_wrap(self): assert mask.shape == (1, 8) # All 8 slots should be valid (we've filled the buffer and window=8=capacity) assert mask[0].sum() == 8 + + +def _simulate_ring_mask(query_len, capacity, offset, window_size): + """Ground truth for concat_window_causal_mask. + + Replays the ring writes for every position before ``offset`` to learn which + absolute position each slot physically holds, then applies the sliding + window predicate directly. Deliberately naive — no modular arithmetic + tricks to mirror the implementation's. + """ + slot_pos = [-1] * capacity + for pos in range(offset): + slot_pos[pos % capacity] = pos + + mask = torch.zeros(query_len, capacity + query_len, dtype=torch.bool) + for i in range(query_len): + q = offset + i + for slot in range(capacity): + k = slot_pos[slot] + mask[i, slot] = k >= 0 and k <= q and (q - k) < window_size + for j in range(query_len): + k = offset + j + mask[i, capacity + j] = k <= q and (q - k) < window_size + return mask + + +def _glimmer_config(window: int = 16, n_layers: int = 4) -> SimpleNamespace: + """All-sliding Muse Glimmer config, so tests isolate the ring path.""" + return SimpleNamespace( + hidden_size=64, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + intermediate_size=128, + num_hidden_layers=n_layers, + vocab_size=128, + rms_norm_eps=1e-5, + sliding_window=window, + layer_types=["sliding_attention"] * n_layers, + layer_rope_theta=[500000.0] * n_layers, + qk_scale_factor=1.0, + max_position_embeddings=256, + ) + + +def _glimmer_caches(model, config, global_len: int = 256): + def zeros(n_layers, seq_len): + return torch.zeros(n_layers, 1, config.num_key_value_heads, seq_len, config.head_dim) + + n_global = max(model.n_global_layers, 1) + n_sliding = model.n_sliding_layers + window = config.sliding_window + return ( + KVCache(zeros(n_global, global_len), zeros(n_global, global_len)), + RingKVCache(zeros(n_sliding, window), zeros(n_sliding, window)), + ) + + +class TestConcatWindowAttention: + """Read-before-write ring path (fetch_and_concat -> attend -> update).""" + + @pytest.mark.parametrize("capacity", [8, 16]) + @pytest.mark.parametrize("window_divisor", [1, 2]) + def test_mask_matches_ring_simulation(self, capacity, window_divisor): + """Mask matches a naive replay of the ring, at every offset through 3 wraps.""" + window_size = capacity // window_divisor + for offset in range(0, 3 * capacity + 3): + for query_len in (1, 2, 3, capacity // 2, capacity): + got = concat_window_causal_mask( + query_len=query_len, + capacity=capacity, + offset=offset, + window_size=window_size, + device="cpu", + ) + want = _simulate_ring_mask(query_len, capacity, offset, window_size) + assert got.shape == (query_len, capacity + query_len) + assert torch.equal(got, want), ( + f"capacity={capacity} window={window_size} " + f"offset={offset} query_len={query_len}" + ) + + def test_mask_from_empty_cache_ignores_ring(self): + """At offset=0 nothing is cached, so only the new-key block is live.""" + mask = concat_window_causal_mask( + query_len=4, capacity=8, offset=0, window_size=8, device="cpu" + ) + assert mask.shape == (4, 12) + assert mask[:, :8].sum() == 0 # ring is entirely unwritten + # New-key block is plain causal + assert torch.equal(mask[:, 8:], torch.tril(torch.ones(4, 4, dtype=torch.bool))) + + def test_mask_excludes_slot_about_to_be_overwritten(self): + """With a full ring, the oldest slot falls outside every query's window.""" + capacity = 8 + mask = concat_window_causal_mask( + query_len=1, capacity=capacity, offset=16, window_size=capacity, device="cpu" + ) + # Query at position 16 attends to 9..16: 7 cached slots + itself. + assert mask[:, :capacity].sum() == capacity - 1 + assert mask[0, capacity] == 1 + # The excluded slot holds position 8 (16 - capacity), written at 8 % 8 = 0. + assert mask[0, 0] == 0 + + def test_fetch_and_concat_does_not_mutate_cache(self): + """fetch_and_concat is read-only; the write happens in update().""" + capacity, n_layers, n_kv, head_dim = 8, 2, 2, 4 + k_buf = torch.randn(n_layers, 1, n_kv, capacity, head_dim) + v_buf = torch.randn(n_layers, 1, n_kv, capacity, head_dim) + cache = RingKVCache(k_buf.clone(), v_buf.clone()) + + k = torch.randn(1, n_kv, 3, head_dim) + v = torch.randn(1, n_kv, 3, head_dim) + full_k, full_v = cache.fetch_and_concat(0, k, v) + + assert full_k.shape == (1, n_kv, capacity + 3, head_dim) + assert torch.equal(cache._k_cache, k_buf), "fetch_and_concat mutated the K cache" + assert torch.equal(cache._v_cache, v_buf), "fetch_and_concat mutated the V cache" + # Concatenated result is [cached window, new keys] + assert torch.equal(full_k[:, :, :capacity], k_buf[0]) + assert torch.equal(full_k[:, :, capacity:], k) + assert torch.equal(full_v[:, :, capacity:], v) + + # The subsequent update must not disturb the already-returned tensors. + snapshot = full_k.clone() + cache.update(0, offset=0, k=k, v=v, query_len=3) + assert torch.equal(full_k, snapshot), "update() aliased fetch_and_concat output" + assert torch.equal(cache._k_cache[0, :, :, :3], k) + + @pytest.mark.parametrize("chunk", [2, 4, 8, 16]) + def test_chunked_prefill_matches_decode_past_window(self, chunk): + """Multi-token prefill matches token-by-token decode, well past the window. + + Regression test for history loss during prefill: a write-first ring lets + a chunk's own K/V evict the oldest ``chunk - 1`` in-window positions + before the chunk's earliest queries read them. Decode (query_len=1) + never hits that, so it is the trusted reference here. + + The guarded model picks the ring path by ``prefill_mode``: Q==1 decode + runs write-first (``prefill_mode=False``), multi-token prefill runs + read-before-write concat (``prefill_mode=True``). Both must match. + + ``chunk`` divides the window so no write straddles the ring boundary. + """ + config = _glimmer_config(window=16) + torch.manual_seed(0) + model = MuseGlimmerModel(config).eval().to(torch.float32) + + seq_len = 48 # 3 full wraps of the 16-slot ring + input_ids = torch.randint(0, config.vocab_size, (1, seq_len)) + + def run(step): + global_cache, sliding_cache = _glimmer_caches(model, config) + outs = [] + for start in range(0, seq_len, step): + end = min(start + step, seq_len) + with torch.no_grad(): + outs.append( + model( + input_ids[:, start:end], + torch.arange(end).unsqueeze(0), + global_cache, + sliding_cache, + prefill_mode=step > 1, + ) + ) + return torch.cat(outs, dim=1) + + reference = run(1) + diff = (run(chunk) - reference).abs().max().item() + assert diff < 1e-4, f"chunk={chunk} diverges from decode reference: {diff}" + + @pytest.mark.parametrize("chunk", [2, 4, 8, 16]) + def test_drafter_chunked_prefill_matches_decode_past_window(self, chunk): + """Same regression check for the drafter, whose layers are all sliding. + + TestRingBufferParity.test_chunked_prefill_matches_single_prefill cannot + catch this: it prefills 48 tokens into a 64-slot ring, so the buffer + never fills and no history is evicted. Here the ring wraps three times. + """ + config = _drafter_config(window=16) + torch.manual_seed(11) + model = DrafterRingModel(config) + model.eval() + + seq_len = 48 # 3 full wraps of the 16-slot ring + input_ids = torch.randint(0, config.vocab_size, (1, seq_len)) + + def run(step): + cache = _make_ring_cache(config) + outs = [] + for start in range(0, seq_len, step): + end = min(start + step, seq_len) + with torch.no_grad(): + outs.append( + model(input_ids[:, start:end], torch.arange(end).unsqueeze(0), cache) + ) + return torch.cat(outs, dim=1) + + reference = run(1) + diff = (run(chunk) - reference).abs().max().item() + assert diff < 1e-4, f"chunk={chunk} diverges from decode reference: {diff}" + + +def _sliding_window_oracle(q, k_log, v_log, offset, window_size): + """Brute-force sliding-window causal attention, independent of the ring. + + For each query ``i`` at absolute position ``p = offset + i``, attend over the + true key/value positions ``[max(0, p - window_size + 1), p]`` taken straight + from a full non-ring K/V log. There is no ring arithmetic here, so a ring + that drops or reorders in-window history cannot hide behind a matching bug in + the reference. + + q, k_log, v_log are (n_heads, seq_len, head_dim); returns (n_heads, Q, head_dim). + """ + n_heads, q_len, head_dim = q.shape + scale = 1.0 / math.sqrt(head_dim) + out = torch.zeros(n_heads, q_len, head_dim) + for i in range(q_len): + p = offset + i + lo = max(0, p - window_size + 1) + ks = k_log[:, lo : p + 1, :] # (n_heads, L, head_dim) + vs = v_log[:, lo : p + 1, :] + scores = torch.einsum("hd,hld->hl", q[:, i, :], ks) * scale + attn = torch.softmax(scores, dim=-1) + out[:, i, :] = torch.einsum("hl,hld->hd", attn, vs) + return out + + +def _ring_chunk_attention(cache, layer_idx, q, k, v, offset, window_size): + """One chunk of sliding attention through the read-before-write ring path. + + Mirrors what the sliding layers do: fetch_and_concat -> attend over the + ``capacity + query_len`` keys with concat_window_causal_mask -> update. q is + (n_heads, Q, head_dim); k, v are (1, n_kv, Q, head_dim). n_heads == n_kv here + to keep the oracle exact. + """ + q_len = q.shape[1] + head_dim = q.shape[-1] + scale = 1.0 / math.sqrt(head_dim) + full_k, full_v = cache.fetch_and_concat(layer_idx, k, v) # (1, n_kv, cap+Q, head_dim) + mask = concat_window_causal_mask( + query_len=q_len, + capacity=cache.capacity(), + offset=offset, + window_size=window_size, + device="cpu", + ) # (Q, cap+Q) + fk = full_k[0] # (n_kv, L, head_dim) + fv = full_v[0] + scores = torch.einsum("hqd,hld->hql", q, fk) * scale # (n_heads, Q, L) + scores = scores.masked_fill(~mask.unsqueeze(0), float("-inf")) + attn = torch.softmax(scores, dim=-1) + out = torch.einsum("hql,hld->hqd", attn, fv) # (n_heads, Q, head_dim) + cache.update(layer_idx, offset, k, v, query_len=q_len) + return out + + +def _ring_chunk_attention_writefirst(cache, layer_idx, q, k, v, offset, window_size): + """One chunk of sliding attention through the write-before-read ring path. + + Mirrors the target's ``main``/decode graph (``prefill_mode=False``): + update_and_fetch writes the chunk into the ring and returns the whole + ``capacity`` window, then attention runs with ring_window_causal_mask. This + is cheaper (no per-step window copy) but only correct for ``query_len == 1``; + a multi-token chunk written after the ring is full evicts the oldest + ``query_len - 1`` in-window positions before its earliest queries read them. + q is (n_heads, Q, head_dim); k, v are (1, n_kv, Q, head_dim). + """ + q_len = q.shape[1] + head_dim = q.shape[-1] + scale = 1.0 / math.sqrt(head_dim) + full_k, full_v = cache.update_and_fetch(layer_idx, offset, k, v, query_len=q_len) + mask = ring_window_causal_mask( + query_len=q_len, + capacity=cache.capacity(), + offset=offset, + window_size=window_size, + device="cpu", + ) # (Q, cap) + fk = full_k[0] # (n_kv, cap, head_dim) + fv = full_v[0] + scores = torch.einsum("hqd,hld->hql", q, fk) * scale # (n_heads, Q, cap) + scores = scores.masked_fill(~mask.unsqueeze(0), float("-inf")) + attn = torch.softmax(scores, dim=-1) + out = torch.einsum("hql,hld->hqd", attn, fv) # (n_heads, Q, head_dim) + return out + + +class TestRingConcatOracleParity: + """Independent brute-force sliding-window oracle for the ring cache primitive. + + TestConcatWindowAttention trusts token-by-token decode as its reference. + This trusts nothing produced by the ring: it recomputes each query's + attention over the true last-``window`` positions from a full K/V log + (_sliding_window_oracle) and compares the ring path against it. The bug the + fix addresses only appears for a multi-token chunk (query_len > 1) written + after the ring is full (offset >= window), so these cases target exactly + that, including a chunk whose read window straddles the physical wrap. + + On the write-first path this test would fail: at query_len > 1 the chunk's + own K/V overwrites the oldest query_len - 1 in-window positions before the + chunk's earliest queries read them. + """ + + WINDOW = 8 + HEAD_DIM = 4 + N_HEADS = 2 # == n_kv, so the oracle stays exact (no GQA broadcast) + + def _run(self, prefill_schedule, chunk_len, seed): + """Prefill with the given chunk sizes, then attend one chunk_len chunk. + + prefill_schedule is the list of chunk sizes used to fill positions + ``[0, sum(schedule))`` into the ring; each must satisfy the no-wrap + write rule. Returns max abs diff between the ring path and the oracle + for the chunk written at offset == sum(schedule). + """ + window, head_dim, n_heads = self.WINDOW, self.HEAD_DIM, self.N_HEADS + offset = sum(prefill_schedule) + total = offset + chunk_len + torch.manual_seed(seed) + k_log = torch.randn(n_heads, total, head_dim) + v_log = torch.randn(n_heads, total, head_dim) + q_log = torch.randn(n_heads, total, head_dim) + + cache = RingKVCache( + torch.zeros(1, 1, n_heads, window, head_dim), + torch.zeros(1, 1, n_heads, window, head_dim), + ) + start = 0 + for size in prefill_schedule: + k = k_log[:, start : start + size, :].unsqueeze(0) + v = v_log[:, start : start + size, :].unsqueeze(0) + cache.update(0, start, k, v, query_len=size) + start += size + + q = q_log[:, offset:total, :] + k = k_log[:, offset:total, :].unsqueeze(0) + v = v_log[:, offset:total, :].unsqueeze(0) + ring_out = _ring_chunk_attention(cache, 0, q, k, v, offset, window) + oracle_out = _sliding_window_oracle(q, k_log, v_log, offset, window) + return (ring_out - oracle_out).abs().max().item() + + def test_decode_control_matches_oracle(self): + """query_len == 1 after the ring is full: passing control on both paths.""" + # Prefill 16 into an 8-slot ring (two full wraps), then a single query. + diff = self._run(prefill_schedule=[8, 8], chunk_len=1, seed=1) + assert diff < 1e-5, f"decode control diverges from oracle: {diff}" + + def test_full_chunk_after_wrap_matches_oracle(self): + """A window-sized chunk written after the ring is full (offset >= window). + + Ring holds positions [8, 15]; the chunk's queries at 16..23 each need + in-window positions the write-first path would have already destroyed. + """ + diff = self._run(prefill_schedule=[8, 8], chunk_len=self.WINDOW, seed=2) + assert diff < 1e-5, f"post-wrap chunk diverges from oracle: {diff}" + + def test_chunk_straddling_physical_wrap_matches_oracle(self): + """A chunk whose read window spans the physical wrap of the ring slots. + + Prefill 8 then 4 leaves the ring physically rotated: slots 0..3 hold + positions 8..11, slots 4..7 hold 4..7. The chunk at offset 12 (Q=4) + writes slots 4..7, and each query's last-8-position window straddles the + slot-0 wrap, so the mask's slot-position reconstruction is exercised. + """ + diff = self._run(prefill_schedule=[8, 4], chunk_len=4, seed=3) + assert diff < 1e-5, f"wrap-straddling chunk diverges from oracle: {diff}" + + @pytest.mark.parametrize("chunk_len", [2, 3, 4]) + def test_partial_chunks_after_wrap_match_oracle(self, chunk_len): + """Multi-token chunks of varied size, all written after the ring fills.""" + # offset = 16 (two wraps), write_start = 0 so chunk_len <= window never wraps. + diff = self._run(prefill_schedule=[8, 8], chunk_len=chunk_len, seed=4 + chunk_len) + assert diff < 1e-5, f"chunk_len={chunk_len} diverges from oracle: {diff}" + + +def _write_index_ancestry_targets(exported_program, write_dim_arg_name): + """Collect the call_function targets feeding a write node's begin index. + + Walks the exported graph backwards from the node whose output name is + ``write_dim_arg_name`` (the ``begin`` bounds of the ring write) and returns + the set of op targets (as strings) that produce it, plus the set of + placeholder names it depends on. + """ + name_to_node = {n.name: n for n in exported_program.graph.nodes} + + def ancestors(name, seen): + if name in seen or name not in name_to_node: + return seen + seen.add(name) + for inp in name_to_node[name].all_input_nodes: + ancestors(inp.name, seen) + return seen + + anc = ancestors(write_dim_arg_name, set()) + targets = {str(name_to_node[a].target) for a in anc if name_to_node[a].op == "call_function"} + placeholders = {a for a in anc if name_to_node[a].op == "placeholder"} + return targets, placeholders + + +class TestConcatWriteIndexIsShapeSymint: + """The ring write index must stay a shape symint, never a runtime data tensor. + + The concat path derives its ring write column from ``offset % capacity`` where + ``offset = seq_len - query_len`` is a shape-derived symint. This must never + collapse into an index read out of tensor *data*: the underlying + slice_update kernel crashes on the GPU backend when its write column comes + from a runtime data tensor rather than from a shape symint. This test pins + that property on the exported graph so a future refactor cannot silently + reintroduce a data-dependent write index. + """ + + def _export_concat_update(self): + capacity, n_kv, head_dim, query_len = 8, 1, 4, 4 + + class _ConcatUpdate(torch.nn.Module): + def forward(self, k, v, k_cache, v_cache, position_ids): + seq_len = position_ids.shape[-1] + q_len = k.shape[-2] + torch._check_is_size(seq_len) + torch._check_is_size(q_len) + offset = seq_len - q_len # shape symint, not a data value + torch._check_is_size(offset) + cache = RingKVCache(k_cache, v_cache) + full_k, full_v = cache.fetch_and_concat(0, k, v) + cache.update(0, offset, k, v, query_len=q_len) + return full_k + 0.0, full_v + 0.0 + + k = torch.zeros(1, n_kv, query_len, head_dim) + v = torch.zeros(1, n_kv, query_len, head_dim) + k_cache = torch.zeros(1, 1, n_kv, capacity, head_dim) + v_cache = torch.zeros(1, 1, n_kv, capacity, head_dim) + position_ids = torch.arange(6, dtype=torch.int32)[None] # seq_len dynamic + + seq = torch.export.Dim("seq", min=query_len, max=capacity * 4) + dynamic_shapes = { + "k": None, + "v": None, + "k_cache": None, + "v_cache": None, + "position_ids": {1: seq}, + } + return torch.export.export( + _ConcatUpdate(), + (k, v, k_cache, v_cache, position_ids), + dynamic_shapes=dynamic_shapes, + ) + + def test_write_index_derives_from_shape_not_data(self): + exported = self._export_concat_update() + + write_nodes = [n for n in exported.graph.nodes if "mutable_slice_update" in str(n.target)] + assert write_nodes, "concat path produced no ring write (mutable_slice_update)" + + for node in write_nodes: + begin_arg = node.args[2] # (x, update, begin, end) + targets, _ = _write_index_ancestry_targets(exported, begin_arg.name) + + # Shape-derived: the write column traces back to a tensor shape query. + assert any("sym_size" in t for t in targets), ( + "ring write index must be derived from a shape symint " + f"(no sym_size op in its ancestry: {sorted(targets)})" + ) + # Not data-derived: nothing pulls a value out of a tensor to index with. + # A data-tensor write index makes the slice_update kernel crash. + data_ops = [t for t in targets if "_local_scalar_dense" in t or ".item" in t] + assert not data_ops, ( + "ring write index must be a shape symint, not a runtime data " + f"tensor, or the slice_update kernel crashes; found {data_ops}" + ) + + +class TestGuardedPathSelection: + """Pin the write-first correctness boundary the Swift guard relies on. + + Muse's decode (``main``) graph is write-before-read: correct at query_len == 1 + but WRONG for query_len > 1 after the ring wraps -- the regime the engine's + guard refuses. (Concat correctness is covered by TestRingConcatOracleParity.) + """ + + WINDOW = 8 + HEAD_DIM = 4 + N_HEADS = 2 # == n_kv, so the oracle stays exact (no GQA broadcast) + + def _prefill_full_ring(self, seed, chunk_len): + """Fill the ring with two full wraps, then stage one chunk at offset=2W.""" + window, head_dim, n_heads = self.WINDOW, self.HEAD_DIM, self.N_HEADS + offset = 2 * window # two full wraps -> ring is full, every slot live + total = offset + chunk_len + torch.manual_seed(seed) + k_log = torch.randn(n_heads, total, head_dim) + v_log = torch.randn(n_heads, total, head_dim) + q_log = torch.randn(n_heads, total, head_dim) + + cache = RingKVCache( + torch.zeros(1, 1, n_heads, window, head_dim), + torch.zeros(1, 1, n_heads, window, head_dim), + ) + for start in (0, window): + k = k_log[:, start : start + window, :].unsqueeze(0) + v = v_log[:, start : start + window, :].unsqueeze(0) + cache.update(0, start, k, v, query_len=window) + + q = q_log[:, offset:total, :] + k = k_log[:, offset:total, :].unsqueeze(0) + v = v_log[:, offset:total, :].unsqueeze(0) + oracle = _sliding_window_oracle(q, k_log, v_log, offset, window) + return cache, q, k, v, offset, oracle + + def test_write_first_matches_oracle_at_decode(self): + """query_len == 1: the write-first path (main graph) is correct.""" + cache, q, k, v, offset, oracle = self._prefill_full_ring(seed=1, chunk_len=1) + out = _ring_chunk_attention_writefirst(cache, 0, q, k, v, offset, self.WINDOW) + diff = (out - oracle).abs().max().item() + assert diff < 1e-5, f"write-first decode diverges from oracle: {diff}" + + @pytest.mark.parametrize("chunk_len", [2, 4, 8]) + def test_write_first_diverges_for_multitoken_after_wrap(self, chunk_len): + """query_len > 1 after wrap: write-first is WRONG (must not match oracle). + + This is the silent corruption the Swift guard exists to prevent: sending + a multi-token chunk through the write-first ``main`` graph evicts + still-needed history. The guard throws instead of producing this. + """ + cache, q, k, v, offset, oracle = self._prefill_full_ring(seed=7, chunk_len=chunk_len) + out = _ring_chunk_attention_writefirst(cache, 0, q, k, v, offset, self.WINDOW) + diff = (out - oracle).abs().max().item() + assert diff > 1e-2, ( + f"write-first should diverge from oracle for query_len={chunk_len} " + f"after wrap, but diff={diff} (bug would be silent)" + ) diff --git a/swift/Sources/CoreAILanguageModels/Handlers/StateHandlerFactory.swift b/swift/Sources/CoreAILanguageModels/Handlers/StateHandlerFactory.swift index a310bba8..e372c51c 100644 --- a/swift/Sources/CoreAILanguageModels/Handlers/StateHandlerFactory.swift +++ b/swift/Sources/CoreAILanguageModels/Handlers/StateHandlerFactory.swift @@ -27,6 +27,8 @@ struct SyncStateHandlerSet { /// True when ALL states are sliding caches (ring buffer model). /// The engine uses compact position_ids [offset, ..., offset+queryLen-1] instead of [0, ..., N-1]. var isAllSlidingCache: Bool + /// True when any state is a sliding (ring) cache (hybrid or pure-ring). + var hasSlidingCache: Bool } /// Creates state handlers from a model's function descriptor. @@ -169,7 +171,8 @@ enum StateHandlerFactory { kvCache: kvCache, additionalStates: additionalStates, hasNonTruncatableStates: hasNonTruncatable, - isAllSlidingCache: classified.allSatisfy { $0.kind == .slidingCache } + isAllSlidingCache: classified.allSatisfy { $0.kind == .slidingCache }, + hasSlidingCache: classified.contains { $0.kind == .slidingCache } ) } } diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift index 8acb75ca..b076d68c 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift @@ -598,6 +598,9 @@ private struct EngineImpl: ~Copyable { var additionalStates: FixedMTLBufferState? var hasNonTruncatableStates: Bool + // True when a sliding-window (ring) KV state is present. + let hasSlidingState: Bool + // Logits — reuses GrowingLogitsBuffer from TensorStorage+CoreAI.swift var logits: GrowingLogitsBuffer @@ -853,6 +856,7 @@ private struct EngineImpl: ~Copyable { self.kvCache = kvCacheLocal self.additionalStates = additionalStatesLocal self.hasNonTruncatableStates = classified.contains(where: { $0.kind == .fixed }) + self.hasSlidingState = classified.contains(where: { $0.kind == .slidingCache }) self.logits = logitsRef self.cachedSampler = nil self.cachedSamplerTemperature = nil @@ -929,6 +933,15 @@ private struct EngineImpl: ~Copyable { let actualTokenCount = tokens.isEmpty ? 1 : tokens.count let queryLength = actualTokenCount + // Write-first `main` graph is correct only at query_len == 1 on a sliding ring. + guard !hasSlidingState || queryLength == 1 else { + throw InferenceRuntimeError.invalidState( + "Sliding-window (ring) model received a \(queryLength)-token decode step on the " + + "write-first 'main' graph, which is only correct for query_len == 1. " + + "Multi-token verify/all-logits over a sliding model needs a Q>1 " + + "concat-with-logits path.") + } + defer { processedTokenCount += actualTokenCount step += 1 @@ -1502,6 +1515,13 @@ private struct EngineImpl: ~Copyable { let actualTokenCount = tokens.isEmpty ? 1 : tokens.count let queryLength = actualTokenCount + // Constrained gen runs the write-first `main` graph: query_len == 1 only. + guard !hasSlidingState || queryLength == 1 else { + throw InferenceRuntimeError.invalidState( + "Sliding-window (ring) model received a \(queryLength)-token constrained step on " + + "the write-first 'main' graph, which is only correct for query_len == 1.") + } + defer { processedTokenCount += actualTokenCount step += 1 @@ -1740,6 +1760,14 @@ private struct EngineImpl: ~Copyable { additionalStates: additionalStates, computeStream: computeStream) } else { + // No prefill graph: write-first `main` chunk is query_len == 1 only. + guard !hasSlidingState || queryLength == 1 else { + throw InferenceRuntimeError.invalidState( + "Sliding-window (ring) model ran a \(queryLength)-token chunk on the " + + "write-first 'main' graph (no prefill graph present), which is only " + + "correct for query_len == 1. Export a prefill graph for multi-token " + + "prompt chunks, or route them through the concat path.") + } let logitsShape = [1, queryLength, vocabSize] try encodeWithStates( function: function, inputs: asyncInputs, diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift index 37dca8c7..2fbd6ca0 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift @@ -45,6 +45,12 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, @u // no logits, so the last prompt token still goes through `function`. private let prefillFunction: InferenceFunction? + // True when a sliding-window (ring) KV state is present. The decode graph + // (`function`) writes the ring write-first, which is only correct for a + // single query token; a multi-token decode step would silently evict + // still-in-window history. The guard in `processTokenBatch` refuses it. + private let hasSlidingState: Bool + // I/O names from descriptor private let logitsName: String @@ -128,6 +134,7 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, @u options: options ) self.options = options + self.hasSlidingState = stateHandlers.hasSlidingCache self.session = GenerationSessionState( kvCache: stateHandlers.kvCache, additionalStates: stateHandlers.additionalStates, @@ -240,6 +247,15 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, @u throw InferenceRuntimeError.invalidState("Cannot process empty token batch") } + // Write-first sliding ring is correct only at query_len == 1; refuse Q>1. + guard !hasSlidingState || batchSize == 1 else { + throw InferenceRuntimeError.invalidState( + "Sliding-window (ring) model received a \(batchSize)-token decode batch on the " + + "write-first 'main' graph, which is only correct for query_len == 1. " + + "Route multi-token steps through the prefill graph; spec-decode verify and " + + "all-logits over a sliding model need a Q>1 concat-with-logits path.") + } + _ = try sessionState.kvCache.ensureCapacity( forContextLength: sessionState.processedTokenCount + batchSize)