From df1fd28b5eeff5859b5c016de78644fef6a43a22 Mon Sep 17 00:00:00 2001
From: lab <l_panos@apple.com>
Date: Fri, 4 Sep 2026 14:07:31 -0700
Subject: [PATCH 1/4] Attend before writing the sliding-window ring, to keep
 prefill history

`RingKVCache.update_and_fetch` wrote the incoming K/V into the ring and
then attended over the result. With `capacity == window_size == W`, a
prefill chunk of `Q` tokens at `offset` writes slots
`[offset%W, offset%W+Q)`, which hold positions `offset-W .. offset-W+Q-1`.
The chunk's first query, at position `offset`, needs positions
`offset-W+1 .. offset`. So the chunk destroys `Q-1` positions of its own
in-window history before its earliest queries ever read them.

Decode is exempt -- at `Q == 1` the only evicted slot is the one that
falls outside the window anyway -- which is why the existing parity
tests never caught this. They exercise decode, or prefill 48 tokens into
a 64-slot ring where the buffer never fills.

Sliding attention now concatenates the new K/V onto the cached window,
attends over `W + Q` keys, and writes to the ring afterwards. Split
`update_and_fetch` into `fetch_and_concat` (read-only) and `update` (the
write); `update_and_fetch` stays as their composition, documented as
decode-only, since nothing needs to change for `Q == 1`.

The mask needed a companion, because `ring_window_causal_mask` describes
the ring *after* the write. `concat_window_causal_mask` describes the
layout before it: `capacity` slots holding positions `< offset`, then
the `Q` new keys. It reconstructs slot positions by walking back from
`offset-1` with the same no-`aten.remainder` trick, and takes the
modulus of `offset + capacity - 1` so the operand stays non-negative at
`offset == 0`, where sympy would otherwise see `Mod(-1, W)`.

Writes still must not wrap: `(offset % W) + Q <= W`. That cannot be
branched on here because `offset` is a symint under `torch.export`, so
it stays a caller obligation -- pick a chunk size dividing the window.

The drafter carries the same fix. Its layers are all sliding, so it was
affected on every layer rather than three in four.

Verified: the new mask is bit-exact against a brute-force replay of the
ring across three wraps for every offset, capacity in {8,16,32}, window
in {capacity, capacity/2}. End-to-end, chunked prefill now matches
token-by-token decode (~3e-06, fp32 noise) where the old path diverged
by 4.6 on the target and 0.22-1.24 on the drafter, growing with chunk
size as expected. Both models still export under `torch.export` with
dynamic shapes, and the exported program matches eager at 0.0 on logits
*and* on the mutated K/V buffers -- which is what confirms
functionalization keeps the read ahead of the write.

The VLM `MuseGlimmerForCausalLMEmbeddings` path is untouched: it forces
`is_sliding = False` on every layer, and its exported graph contains no
ring ops at all.

(cherry picked from commit 3509a1a56b2275e0b6167638a1d4694697d3cf85)
---
 .../models/macos/muse_glimmer.py              |  22 +-
 .../models/macos/muse_glimmer_drafter_ring.py |  25 ++-
 .../coreai_models/primitives/macos/cache.py   | 130 ++++++++++-
 .../test_ring_buffer_parity.py                | 209 +++++++++++++++++-
 4 files changed, 364 insertions(+), 22 deletions(-)

diff --git a/python/src/coreai_models/models/macos/muse_glimmer.py b/python/src/coreai_models/models/macos/muse_glimmer.py
index 95c4585b..9df34248 100644
--- a/python/src/coreai_models/models/macos/muse_glimmer.py
+++ b/python/src/coreai_models/models/macos/muse_glimmer.py
@@ -18,6 +18,11 @@
 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 read-before-write: each step concatenates the incoming
+K/V onto the cached window, attends over ``window_size + query_len`` keys, and
+only then writes into the ring. Writing first would let a multi-token prefill
+chunk evict window history that its own earliest queries still need.
 """
 
 import gc
@@ -44,7 +49,7 @@
     _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
 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
@@ -132,10 +137,17 @@ 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
+            # Read-before-write: attend over [cached window ++ new K/V], then
+            # write. Writing first would let this chunk's own K/V evict the
+            # oldest query_len - 1 in-window positions that its earliest
+            # queries still need.
+            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
             )
-            attn_output = self.sdpa(query=query, key=key, value=value, attn_mask=sliding_mask)
+            sliding_kv_cache.update(self.cache_slot_idx, offset, key, value, query_len=query_len)
         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
@@ -235,7 +247,7 @@ def forward(
 
         sliding_mask = None
         if sliding_kv_cache is not None:
-            sliding_mask = ring_window_causal_mask(
+            sliding_mask = concat_window_causal_mask(
                 query_len,
                 sliding_kv_cache.capacity(),
                 offset,
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..d8c44dad 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,11 @@
 - 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 is read-before-write: each step concatenates the incoming
+K/V onto the cached window, attends over ``window_size + query_len`` keys, and
+only then writes into the ring. Writing first would let a multi-token prefill
+chunk evict window history that its own earliest queries still need.
 """
 
 import gc
@@ -38,7 +43,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 +116,23 @@ 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: attend over [cached window ++ new K/V], then
+            # write. Writing first would let this chunk's own K/V evict the
+            # oldest query_len - 1 in-window positions that its earliest
+            # queries still need.
+            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 +188,7 @@ 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(
+            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..de142503 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,80 @@ 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 window ++ new keys]``.
+
+    Companion to :meth:`RingKVCache.fetch_and_concat`. Unlike
+    :func:`ring_window_causal_mask` -- which describes the ring *after* the new
+    K/V has been written into it -- this describes the key layout *before* the
+    write: ``capacity`` ring slots holding positions ``< offset``, followed by
+    the ``query_len`` new keys at positions ``offset .. offset + query_len - 1``.
+
+    Reading the ring before overwriting it is what makes a multi-token step
+    correct. With a write-first ring the chunk's own K/V evicts the oldest
+    ``query_len - 1`` positions of the window, which the earliest queries in the
+    same chunk still need to attend to.
+
+    Args:
+        query_len: number of query positions in this forward call
+        capacity: ring buffer size (number of physical slots)
+        offset: absolute position of the first query token
+        window_size: sliding window size (typically == capacity)
+        device: target device
+
+    Returns:
+        Boolean mask [query_len, capacity + query_len] — True means "attend
+        to this key". Columns ``[0, capacity)`` index ring slots, columns
+        ``[capacity, capacity + query_len)`` index the new keys in order.
+    """
+    row = torch.arange(query_len, device=device)
+    q_pos = row.unsqueeze(-1) + offset  # (query_len, 1)
+
+    # --- cached ring slots -------------------------------------------------
+    # The newest position already in the ring is offset - 1. Reconstruct each
+    # slot's absolute position by walking backwards from it, same trick as
+    # ring_window_causal_mask (no tensor aten.remainder). Adding `capacity`
+    # keeps the modulo operand non-negative when 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)
+
+    # Causality is implicit here: every cached position is < offset <= q_pos.
+    # Slots not yet written reconstruct to a negative position — drop those.
+    ring_mask = ((q_pos - k_pos) < window_size) & (k_pos >= 0)
+
+    # --- new keys ----------------------------------------------------------
+    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 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.
+    at position % capacity and never grows.
+
+    Two usage patterns:
+
+    * ``fetch_and_concat`` -> attention -> ``update`` (read-before-write).
+      Attends over ``capacity + query_len`` keys with
+      :func:`concat_window_causal_mask`. Correct for any ``query_len``, since
+      the incoming K/V cannot evict history the same step still needs.
+    * ``update_and_fetch`` -> attention (write-before-read). Attends over
+      ``capacity`` keys with :func:`ring_window_causal_mask`. Cheaper, but only
+      correct for ``query_len == 1``: a multi-token chunk overwrites the oldest
+      ``query_len - 1`` in-window positions before reading them.
     """
 
     def __init__(self, k_cache: torch.Tensor, v_cache: torch.Tensor) -> None:
@@ -343,21 +414,43 @@ 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,
+        k: torch.Tensor,
+        v: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Return this layer's cached window with ``k``/``v`` appended.
+
+        Does not mutate the cache — call :meth:`update` after attention. The
+        returned tensors have seq length ``capacity + k.size(-2)`` and pair with
+        a :func:`concat_window_causal_mask` of matching width.
+        """
+        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)
+        # torch.cat copies, so the results do not alias the cache and stay
+        # valid across the subsequent update().
+        return torch.cat([cached_k, k], dim=-2), torch.cat([cached_v, v], dim=-2)
+
+    def update(
         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.
+    ) -> None:
+        """Write K/V into this layer's ring slots starting at ``offset % capacity``.
 
         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.
+        ``(offset % capacity) + query_len <= capacity`` must hold. Decode
+        (query_len=1) always satisfies this. For chunked prefill, choose a chunk
+        size that divides ``capacity`` (e.g. ``capacity // 2``) so that no chunk
+        straddles the wrap point. The bound cannot be branched on here because
+        ``offset`` is a symint under ``torch.export``.
 
         Raises:
             RuntimeError: If the write would wrap around the ring buffer.
@@ -401,6 +494,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..b9659832 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,6 +8,10 @@
 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.
 """
 
 from types import SimpleNamespace
@@ -15,8 +19,14 @@
 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 +259,200 @@ 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.
+
+        ``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,
+                        )
+                    )
+            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}"

From 46d0445e33e5f079d08dc06ec05f01d76711d057 Mon Sep 17 00:00:00 2001
From: sukru tikves <sukrutikves@gmail.com>
Date: Sun, 13 Sep 2026 22:17:10 -0700
Subject: [PATCH 2/4] Add Q>1 post-wrap parity test for sliding-window ring
 cache

The concat parity tests trust token-by-token decode as their reference.
This adds an independent brute-force sliding-window oracle that recomputes
each query's attention over the true last-W key/value positions from a full
non-ring K/V log, and compares the fetch_and_concat ring path against it.

Covers a multi-token chunk (Q>1) written after the ring is full
(offset >= window), including a chunk whose read window straddles the
physical wrap, plus a Q=1 decode control. On the old write-first path these
diverge from the oracle by 0.49-1.41; the read-before-write path matches
within fp32 noise.

(cherry picked from commit d8d13ff6f0eed40d82fa590cddd9de7f127fa1a4)
---
 .../test_ring_buffer_parity.py                | 143 ++++++++++++++++++
 1 file changed, 143 insertions(+)

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 b9659832..770a91a9 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
@@ -14,6 +14,7 @@
 update_and_fetch) is correct for multi-token prefill chunks.
 """
 
+import math
 from types import SimpleNamespace
 
 import pytest
@@ -456,3 +457,145 @@ def run(step):
         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
+
+
+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}"

From fed102d635e245c99702de69d2ff4b7d0aad099f Mon Sep 17 00:00:00 2001
From: sukru tikves <sukrutikves@gmail.com>
Date: Mon, 28 Sep 2026 16:30:17 -0700
Subject: [PATCH 3/4] Guard ring concat write index stays a shape symint

The concat prefill path writes into the ring at offset % capacity, where
offset = seq_len - query_len is derived from tensor shapes. That write column
must never come from tensor data: the slice_update kernel crashes on the GPU
backend when its write index is a runtime data tensor instead of a shape
symint.

Add an export-based test that pins the property: it exports the
fetch_and_concat -> update path with a dynamic sequence length, finds the ring
write node, and asserts its begin index traces back to a shape query
(sym_size) and never to a value pulled out of a tensor. A future refactor that
reintroduced a data-dependent write index would flip this test red before it
reached a device.

(cherry picked from commit 1c712bccac124a31aabfc91a1c0ba88b00ebd8e1)
---
 .../test_ring_buffer_parity.py                | 96 +++++++++++++++++++
 1 file changed, 96 insertions(+)

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 770a91a9..9ac3b819 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
@@ -599,3 +599,99 @@ def test_partial_chunks_after_wrap_match_oracle(self, chunk_len):
         # 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}"
+            )

From 2f2b50265785f6e50ff9dfbecb54c99b5be1687f Mon Sep 17 00:00:00 2001
From: Sukru Tikves <sukrutikves@gmail.com>
Date: Sat, 10 Oct 2026 17:10:07 -0700
Subject: [PATCH 4/4] Guard sliding-window ring: concat on prefill, fast
 write-first on Q=1 decode

Gate the sliding-window ring access on prefill_mode: the prefill graph reads the
ring before writing (concat), correct for any chunk after wrap; the main/decode
graph keeps the fast write-first path, correct for Q==1. Add a fail-loud Swift
guard so a Q>1 call on the write-first main path (spec-decode verify, PPL) throws
instead of silently corrupting, keyed on the presence of a sliding-ring state.
---
 .../models/macos/muse_glimmer.py              | 87 +++++++++++++----
 .../macos/muse_glimmer_drafter_dflash.py      |  3 +-
 .../models/macos/muse_glimmer_drafter_ring.py | 13 +--
 .../coreai_models/primitives/macos/cache.py   | 82 +++++-----------
 .../test_ring_buffer_parity.py                | 97 +++++++++++++++++++
 .../Handlers/StateHandlerFactory.swift        |  5 +-
 .../CoreAIPipelinedEngine.swift               | 28 ++++++
 .../CoreAISequentialEngine.swift              | 16 +++
 8 files changed, 241 insertions(+), 90 deletions(-)

diff --git a/python/src/coreai_models/models/macos/muse_glimmer.py b/python/src/coreai_models/models/macos/muse_glimmer.py
index 9df34248..84eec242 100644
--- a/python/src/coreai_models/models/macos/muse_glimmer.py
+++ b/python/src/coreai_models/models/macos/muse_glimmer.py
@@ -19,10 +19,9 @@
 Sliding layers use RingKVCache at window_size capacity; global layers use
 standard KVCache with dynamic sequence length.
 
-Sliding attention is read-before-write: each step concatenates the incoming
-K/V onto the cached window, attends over ``window_size + query_len`` keys, and
-only then writes into the ring. Writing first would let a multi-token prefill
-chunk evict window history that its own earliest queries still need.
+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
@@ -49,13 +48,38 @@
     _load_tensors_for_keys,
     _resolve_safetensors_files,
 )
-from coreai_models.primitives.macos.cache import KVCache, RingKVCache, concat_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__()
@@ -104,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
@@ -137,17 +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:
-            # Read-before-write: attend over [cached window ++ new K/V], then
-            # write. Writing first would let this chunk's own K/V evict the
-            # oldest query_len - 1 in-window positions that its earliest
-            # queries still need.
-            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)
+            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
@@ -185,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),
@@ -194,6 +230,7 @@ def forward(
             global_cache,
             sliding_kv_cache,
             sliding_mask,
+            prefill_mode,
         )
         r = self.post_attention_layernorm(r)
         h = x + r
@@ -234,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]
@@ -247,7 +285,8 @@ def forward(
 
         sliding_mask = None
         if sliding_kv_cache is not None:
-            sliding_mask = concat_window_causal_mask(
+            sliding_mask = _sliding_window_mask(
+                prefill_mode,
                 query_len,
                 sliding_kv_cache.capacity(),
                 offset,
@@ -264,6 +303,7 @@ def forward(
                 global_cache,
                 sliding_kv_cache,
                 sliding_mask,
+                prefill_mode,
             )
         h = self.norm(h)
         h = h * self.output_multiplier
@@ -423,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)
@@ -589,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]
@@ -602,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,
@@ -621,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)
@@ -660,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 d8c44dad..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
@@ -14,10 +14,8 @@
 - Shares embed_tokens and lm_head with the target model
 - Ring buffer KV cache with explicit attention mask
 
-Sliding attention is read-before-write: each step concatenates the incoming
-K/V onto the cached window, attends over ``window_size + query_len`` keys, and
-only then writes into the ring. Writing first would let a multi-token prefill
-chunk evict window history that its own earliest queries still need.
+Sliding attention reads before writing (concat), so a multi-token chunk can't
+evict window history it still needs.
 """
 
 import gc
@@ -116,10 +114,7 @@ def forward(
         key = self.rope(key, position_ids=rope_positions, freqs=freqs)
 
         if cache is not None:
-            # Read-before-write: attend over [cached window ++ new K/V], then
-            # write. Writing first would let this chunk's own K/V evict the
-            # oldest query_len - 1 in-window positions that its earliest
-            # queries still need.
+            # 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
@@ -188,6 +183,8 @@ def forward(
         h = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + self.config.rms_norm_eps)
 
         if cache is not None:
+            # 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(),
diff --git a/python/src/coreai_models/primitives/macos/cache.py b/python/src/coreai_models/primitives/macos/cache.py
index de142503..848b2049 100644
--- a/python/src/coreai_models/primitives/macos/cache.py
+++ b/python/src/coreai_models/primitives/macos/cache.py
@@ -338,39 +338,20 @@ def concat_window_causal_mask(
     window_size: int,
     device: torch.device,
 ) -> torch.Tensor:
-    """Sliding-window causal mask for ``[cached ring window ++ new keys]``.
+    """Sliding-window causal mask for ``[cached ring ++ new keys]``, before the write.
 
-    Companion to :meth:`RingKVCache.fetch_and_concat`. Unlike
-    :func:`ring_window_causal_mask` -- which describes the ring *after* the new
-    K/V has been written into it -- this describes the key layout *before* the
-    write: ``capacity`` ring slots holding positions ``< offset``, followed by
-    the ``query_len`` new keys at positions ``offset .. offset + query_len - 1``.
+    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.
 
-    Reading the ring before overwriting it is what makes a multi-token step
-    correct. With a write-first ring the chunk's own K/V evicts the oldest
-    ``query_len - 1`` positions of the window, which the earliest queries in the
-    same chunk still need to attend to.
-
-    Args:
-        query_len: number of query positions in this forward call
-        capacity: ring buffer size (number of physical slots)
-        offset: absolute position of the first query token
-        window_size: sliding window size (typically == capacity)
-        device: target device
-
-    Returns:
-        Boolean mask [query_len, capacity + query_len] — True means "attend
-        to this key". Columns ``[0, capacity)`` index ring slots, columns
-        ``[capacity, capacity + query_len)`` index the new keys in order.
+    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)
 
-    # --- cached ring slots -------------------------------------------------
-    # The newest position already in the ring is offset - 1. Reconstruct each
-    # slot's absolute position by walking backwards from it, same trick as
-    # ring_window_causal_mask (no tensor aten.remainder). Adding `capacity`
-    # keeps the modulo operand non-negative when offset == 0.
+    # 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
@@ -378,11 +359,9 @@ def concat_window_causal_mask(
     ring_back = torch.where(diff >= 0, diff, diff + capacity)
     k_pos = (last_pos - ring_back).unsqueeze(0)  # (1, capacity)
 
-    # Causality is implicit here: every cached position is < offset <= q_pos.
-    # Slots not yet written reconstruct to a negative position — drop those.
+    # Cached positions are all < offset, so causality is implicit; drop unwritten slots.
     ring_mask = ((q_pos - k_pos) < window_size) & (k_pos >= 0)
 
-    # --- new keys ----------------------------------------------------------
     new_pos = (row + offset).unsqueeze(0)  # (1, query_len)
     new_mask = (new_pos <= q_pos) & ((q_pos - new_pos) < window_size)
 
@@ -390,21 +369,12 @@ def concat_window_causal_mask(
 
 
 class RingKVCache:
-    """Ring-buffer KV cache for sliding window attention layers.
-
-    Fixed-size cache [n_layers, 1, n_kv_heads, capacity, head_dim] that writes
-    at position % capacity and never grows.
-
-    Two usage patterns:
+    """Fixed-size ring KV cache for sliding-window layers; writes at position % capacity.
 
-    * ``fetch_and_concat`` -> attention -> ``update`` (read-before-write).
-      Attends over ``capacity + query_len`` keys with
-      :func:`concat_window_causal_mask`. Correct for any ``query_len``, since
-      the incoming K/V cannot evict history the same step still needs.
-    * ``update_and_fetch`` -> attention (write-before-read). Attends over
-      ``capacity`` keys with :func:`ring_window_causal_mask`. Cheaper, but only
-      correct for ``query_len == 1``: a multi-token chunk overwrites the oldest
-      ``query_len - 1`` in-window positions before reading them.
+    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:
@@ -420,19 +390,15 @@ def fetch_and_concat(
         k: torch.Tensor,
         v: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Return this layer's cached window with ``k``/``v`` appended.
-
-        Does not mutate the cache — call :meth:`update` after attention. The
-        returned tensors have seq length ``capacity + k.size(-2)`` and pair with
-        a :func:`concat_window_causal_mask` of matching width.
+        """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)
-        # torch.cat copies, so the results do not alias the cache and stay
-        # valid across the subsequent update().
+        # 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)
 
     def update(
@@ -443,17 +409,13 @@ def update(
         v: torch.Tensor,
         query_len: int | None = None,
     ) -> None:
-        """Write K/V into this layer's ring slots starting at ``offset % capacity``.
-
-        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 a chunk
-        size that divides ``capacity`` (e.g. ``capacity // 2``) so that no chunk
-        straddles the wrap point. The bound cannot be branched on here because
-        ``offset`` is a symint under ``torch.export``.
+        """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]
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 9ac3b819..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
@@ -398,6 +398,10 @@ def test_chunked_prefill_matches_decode_past_window(self, chunk):
         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)
@@ -419,6 +423,7 @@ def run(step):
                             torch.arange(end).unsqueeze(0),
                             global_cache,
                             sliding_cache,
+                            prefill_mode=step > 1,
                         )
                     )
             return torch.cat(outs, dim=1)
@@ -513,6 +518,37 @@ def _ring_chunk_attention(cache, layer_idx, q, k, v, offset, window_size):
     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.
 
@@ -695,3 +731,64 @@ def test_write_index_derives_from_shape_not_data(self):
                 "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)