fix: make write_token dynamic (#29271)

python/sglang/srt/hardware_backend/mlx/kv_cache/attention_kv_cache.py

@@ -112,9 +112,12 @@ class ContiguousAttentionKVCache:
 
     def write_token(self, k: mx.array, v: mx.array) -> None:
         """Write one token. k, v shape: (1, n_kv_heads, 1, head_dim)."""
-        self.keys[:, :, self.offset : self.offset + 1, :] = k
-        self.values[:, :, self.offset : self.offset + 1, :] = v
-        self.offset += 1
+        end = self.offset + 1
+        if end > self.max_seq_len:
+            self._grow(end)
+        self.keys[:, :, self.offset : end, :] = k
+        self.values[:, :, self.offset : end, :] = v
+        self.offset = end
 
     def get_kv(self) -> tuple[mx.array, mx.array]:
         """Return valid K/V: (1, n_kv_heads, offset, head_dim)."""

python/sglang/srt/hardware_backend/mlx/model_runner.py

@@ -1198,10 +1198,6 @@ class MlxModelRunner:
         """
         caches = prev.caches
 
-        # TODO (changminbark): Need to fix
-        # ContiguousAttentionKVCache.write_token to accommodate dynamic growing
-        # like ContiguousAttentionKVCache.update_and_fetch.
-
         # After prev's graph ran, each attention KV cache offset was
         # bumped by one per layer - attention wrapper's `write_token`
         # mutates the Python offset synchronously at graph-build time.

test/registered/unit/hardware_backend/mlx/test_attention_patching.py

@@ -190,6 +190,31 @@ class TestMlxAttentionPatching(unittest.TestCase):
         self.assertEqual(out.shape, (1, 1, 4))
         self.assertEqual(inner.o_proj.last_input_shape, (1, 1, 4))
 
+    def test_write_token_grows_buffer_past_max_seq_len(self):
+        max_seq_len = 4
+        cache = ContiguousAttentionKVCache(
+            n_kv_heads=1, head_dim=2, max_seq_len=max_seq_len, dtype=mx.float32
+        )
+        n_tokens = max_seq_len * 2 + 1  # force at least one grow past the boundary
+
+        for t in range(n_tokens):
+            k = mx.full((1, 1, 1, 2), t, dtype=mx.float32)
+            v = mx.full((1, 1, 1, 2), -t, dtype=mx.float32)
+            cache.write_token(k, v)
+
+        self.assertEqual(cache.offset, n_tokens)
+        self.assertGreaterEqual(cache.max_seq_len, n_tokens)
+
+        keys, values = cache.get_kv()
+        mx.eval(keys, values)
+        self.assertEqual(keys.shape, (1, 1, n_tokens, 2))
+        self.assertEqual(values.shape, (1, 1, n_tokens, 2))
+        # Every token (including those written before the grow) is preserved
+        # at its original position.
+        for t in range(n_tokens):
+            self.assertEqual(keys[0, 0, t, 0].item(), float(t))
+            self.assertEqual(values[0, 0, t, 0].item(), float(-t))
+
     def test_attn_config_uses_float_dtype_for_quantized_projection(self):
         runner = object.__new__(MlxModelRunner)
         attn = FakeAttention()