Append DeepSeek V4.1 window snapshots during cache writes

Combine immutable window-prefix copies with appended KV encoding during verify and draft preparation. Preserve lockfile platform markers.

Cherry-picked from cb5b34e and 02d7328 (PR #108), with dependency conflicts resolved against local main. Validation: 5 Ignition snapshot preparation tests and 27 FlashInfer snapshot/word-exchange contract tests passed. CUDA and full-model export checks were not run; target GPUs are occupied.

Design-Signed-Off-By: Codex GPT-6; design-review=1c01

Author: lukealonso

ignition/components/compressed_mla.py

@@ -24,34 +24,51 @@ def validate_mla_compute_mode(mode):
 
 
 @torch.library.custom_op("flare::compressed_cache_write", mutates_args=("cache",))
-def _write(kv: torch.Tensor, cache: torch.Tensor, slots: torch.Tensor, handle: int) -> None:
-    mla.write_cache(kv, cache, slots, plan=plan_from_handle(handle))
+def _write(kv: torch.Tensor, cache: torch.Tensor, slots: torch.Tensor, handle: int,
+           history: torch.Tensor | None = None, history_rows: torch.Tensor | None = None) -> None:
+    mla.write_cache(kv, cache, slots, plan=plan_from_handle(handle), history=history, history_rows=history_rows)
 
 
 @_write.register_fake
-def _write_fake(kv, cache, slots, handle):
+def _write_fake(kv, cache, slots, handle, history=None, history_rows=None):
     pass
 
 
 class PreparedCompressedCacheWriter:
-    def __init__(self, values, cache, slots, *, page_size, kind, session, name, cache_format="deepseek_v41"):
+    def __init__(self, values, cache, slots, *, page_size, kind, session, name, cache_format="deepseek_v41",
+                 history=None, history_rows=None, snapshot_tokens=0):
+        if bool(snapshot_tokens) != (history is not None and history_rows is not None):
+            raise ValueError("snapshot cache writers require history, row IDs, and a prefix extent")
+        if (history is None) != (history_rows is None):
+            raise ValueError("history and history_rows must be supplied together")
         self.plan = mla.plan_cache_writer(mla.CacheWriterQuery(max_rows=values.shape[0],
             page_size=page_size, cache_kind=kind, cache_format=cache_format,
-            slot_dtype=str(slots.dtype).removeprefix("torch.")),
+            slot_dtype=str(slots.dtype).removeprefix("torch."), snapshot_tokens=snapshot_tokens,
+            max_requests=0 if history_rows is None else history_rows.numel()),
             device=values.device)
         self.cache, self.slots = cache, slots
+        self.history, self.history_rows = history, history_rows
 
         def prepare(state):
             source = torch.zeros_like(values)
-            target = torch.empty_strided((1, cache.shape[1]), cache.stride(), dtype=cache.dtype, device=cache.device)
+            rows = self.plan.query.max_requests or 1
+            target = torch.empty_strided((rows, cache.shape[1]), cache.stride(), dtype=cache.dtype, device=cache.device)
             positions = torch.arange(values.shape[0], dtype=slots.dtype, device=slots.device)
-            positions.masked_fill_(positions >= page_size, -1)
-            return PreparedCall(run=lambda: state.run(source, target, positions))
+            if snapshot_tokens:
+                suffix = page_size - snapshot_tokens
+                positions = positions // suffix * page_size + snapshot_tokens + positions % suffix
+                positions.masked_fill_(positions >= rows * page_size, -1)
+                seed = torch.zeros((rows + 1, snapshot_tokens * 528), dtype=torch.uint8, device=cache.device)
+                selected = torch.arange(1, rows + 1, dtype=torch.int64, device=cache.device)
+            else:
+                positions.masked_fill_(positions >= page_size, -1)
+                seed = selected = None
+            return PreparedCall(run=lambda: state.run(source, target, positions, history=seed, history_rows=selected))
 
         declare(session, (self.plan.request(name=name, prepare_call=prepare),))
 
     def __call__(self, values):
-        _write(values, self.cache, self.slots, self.plan.handle)
+        _write(values, self.cache, self.slots, self.plan.handle, self.history, self.history_rows)
 
 
 @torch.library.custom_op("flare::compressed_mla", mutates_args=("out", "scratch"))

ignition/recipes/deepseek41/attention.py

@@ -323,7 +323,10 @@ def zeros(shape, dtype):
             weights["output"], scratch.output, session=session, name=name + ".output", workspace=workspace)
         self.window_writer = PreparedCompressedCacheWriter(scratch.window_values, self.window_cache,
             self.window_slots, page_size=self.window_page_size, kind="swa", session=session,
-            name=f"{name}.window_write")
+            name=f"{name}.window_write",
+            history=self.window if mode in ("verify", "draft") else None,
+            history_rows=metadata.state_rows if mode in ("verify", "draft") else None,
+            snapshot_tokens=128 if mode in ("verify", "draft") else 0)
         self.global_kv = None
         self.selectors = []
         if "compressor" in weights:
@@ -396,7 +399,7 @@ def __call__(self, source):
         if self.owns_window_metadata:
             self.window_metadata()
         position = self.window_metadata.position
-        if self.mode != "decode":
+        if self.mode == "prefill":
             gather_rows(self.window, m.state_rows, self.window_cache[:, :128 * 528])
         self.projections["qkv_a"](source)
         if self.norm_quant is None:

pyproject.toml

@@ -22,7 +22,7 @@ prepare = [
     "torch==2.14.1+cu132",
     "torchvision==0.29.1+cu132",
     "transformers[torch,vision,chat_template] @ git+https://github.com/huggingface/transformers.git@1d9654e78a23fb167ef539cbd699dfbd3e4830cf",
-    "flashinfer-python[b12x] @ git+https://github.com/local-inference-lab/flashinfer.git@f91f257e079d8fa632d64cc6db49a1e609ff02ed",
+    "flashinfer-python[b12x] @ git+https://github.com/local-inference-lab/flashinfer.git@d80824f62e6509a5eb9247d9f3983263da01dfdd",
     "cuda-python==13.3.1",
     "cuda-tile[tileiras]==1.6.0",
     "safetensors==0.8.0",

tests/components/test_compressed_cache_snapshot.py

@@ -0,0 +1,53 @@
+"""Preparation declares immutable snapshot operands and bounded priming storage."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from ignition.components import compressed_mla
+
+
+@pytest.mark.parametrize("requests,width", [(1, 8), (3, 7)])
+def test_snapshot_writer_declares_operands_and_primes_valid_private_pages(monkeypatch, requests, width):
+    source = torch.zeros(requests * width, 512, dtype=torch.bfloat16)
+    cache = torch.empty(requests, (128 + width) * 528, dtype=torch.uint8)
+    history = torch.empty(9, 128 * 528, dtype=torch.uint8)
+    rows = torch.arange(requests, dtype=torch.int64)
+    slots = torch.empty(requests * width, dtype=torch.int64)
+    requests_out, calls = [], []
+
+    def plan(query, *, device):
+        return SimpleNamespace(query=query, handle=1,
+            request=lambda **kwargs: SimpleNamespace(**kwargs))
+
+    monkeypatch.setattr(compressed_mla.mla, "plan_cache_writer", plan)
+    monkeypatch.setattr(compressed_mla, "declare", lambda session, values: requests_out.extend(values))
+    monkeypatch.setattr(compressed_mla, "_write", lambda *args: calls.append(args))
+    writer = compressed_mla.PreparedCompressedCacheWriter(source, cache, slots,
+        page_size=128 + width, kind="swa", session=object(), name="window",
+        history=history, history_rows=rows, snapshot_tokens=128)
+    assert writer.plan.query.snapshot_tokens == 128
+    assert writer.plan.query.max_requests == requests
+    observed = []
+    call = requests_out[0].prepare_call(SimpleNamespace(run=lambda *args, **kwargs: observed.append((args, kwargs))))
+    call.invoke()
+    args, kwargs = observed[0]
+    assert args[1].shape == cache.shape
+    assert kwargs["history"].shape == (requests + 1, 128 * 528)
+    assert torch.equal(kwargs["history_rows"], torch.arange(1, requests + 1))
+    assert torch.all(args[2] % (128 + width) >= 128)
+    writer(source)
+    assert calls[0][4] is history and calls[0][5] is rows
+
+
+@pytest.mark.parametrize("history,rows,tokens", [(None, None, 128), (True, None, 128), (True, True, 0)])
+def test_snapshot_writer_rejects_incomplete_operands(history, rows, tokens):
+    source = torch.empty(8, 512, dtype=torch.bfloat16)
+    cache = torch.empty(1, 136 * 528, dtype=torch.uint8)
+    selected = torch.zeros(1, dtype=torch.int64)
+    with pytest.raises(ValueError):
+        compressed_mla.PreparedCompressedCacheWriter(source, cache, torch.empty(8, dtype=torch.int64),
+            page_size=136, kind="swa", session=object(), name="window",
+            history=cache if history else None, history_rows=selected if rows else None,
+            snapshot_tokens=tokens)

tests/recipes/test_deepseek41_snapshot_writer_cuda.py

@@ -0,0 +1,131 @@
+"""Compiled window snapshots preserve immutable history and accepted prefixes."""
+
+import os
+
+import pytest
+import torch
+from b12x.preparation import PreparationSession
+from cuda.bindings import driver as cuda
+
+from ignition.components.compressed_mla import PreparedCompressedCacheWriter
+from ignition.components.paged_batch import Metadata
+from ignition.components.ring_commit import PreparedRingCommit
+from ignition.components.state import StateLayout, StateStorage, gather_rows
+from ignition.components.window_indices import PreparedWindowIndices
+from ignition.preparation.compilation import compile_forward
+from ignition.preparation.graphs import PreparedGraphBundle
+
+pytestmark = pytest.mark.skipif(os.environ.get("FLARE_RUN_CUDA_TESTS") != "1",
+                              reason="requires an explicitly assigned CUDA device")
+
+
+class SnapshotCase:
+    def __init__(self, requests, width, mode, session, *, fused):
+        device = torch.device("cuda:0")
+        self.mode, self.fused = mode, fused
+        layout = StateLayout.from_specs((("before", (256,), torch.uint8),
+            ("window", (128 * 528,), torch.uint8), ("after", (256,), torch.uint8)))
+        self.storage = StateStorage.allocate(layout, requests + 3, device)
+        self.history = self.storage.region("window")
+        self.metadata = Metadata.allocate(requests * width, requests, 1, device,
+                                         max_query_tokens=width)
+        self.indices = PreparedWindowIndices(self.metadata, mode=mode,
+            tokens=requests * width, page_size=128 + width,
+            width=128 + width if mode == "draft" else 128)
+        self.values = torch.zeros((requests * width, 512), dtype=torch.bfloat16, device=device)
+        self.counts = torch.zeros(requests, dtype=torch.int32, device=device)
+        self.commit = PreparedRingCommit(self.storage.backing,
+            (layout.region("window").offset,), self.metadata.state_rows, self.metadata.starts,
+            token_capacity=width, window_tokens=128, record_bytes=528)
+        self.cache = self.commit.staging[0]
+        self.observed = torch.empty_like(self.cache)
+        self.writer = PreparedCompressedCacheWriter(self.values, self.cache, self.indices.window_slots,
+            page_size=128 + width, kind="swa", session=session, name=f"window.{fused}",
+            history=self.history if fused else None,
+            history_rows=self.metadata.state_rows if fused else None,
+            snapshot_tokens=128 if fused else 0)
+
+    def __call__(self):
+        self.indices()
+        if not self.fused:
+            gather_rows(self.history, self.metadata.state_rows, self.cache[:, :128 * 528])
+        self.writer(self.values)
+        self.observed.copy_(self.cache)
+        if self.mode == "verify":
+            self.commit(self.counts)
+        # Consumers and accepted commits must finish before private inputs are reused.
+        self.values.fill_(37)
+        self.cache.fill_(129)
+
+
+@pytest.mark.parametrize("requests", [1, 3])
+@pytest.mark.parametrize("mode,width", [("verify", 8), ("draft", 7)])
+@torch.no_grad()
+def test_compiled_native_snapshot_writer_preserves_ragged_state_and_lifetimes(requests, mode, width):
+    device = torch.device("cuda:0")
+    with PreparationSession(device=device, autotune=False, compile_workers=2) as session:
+        cases = [SnapshotCase(requests, width, mode, session, fused=fused) for fused in (False, True)]
+        compiled = [compile_forward(case) for case in cases]
+        graphs = {}
+        counts = []
+        for name, forward in zip(("separate", "fused"), compiled, strict=True):
+            graph = torch.cuda.CUDAGraph(keep_graph=True)
+            with session.capture(), torch.cuda.graph(graph):
+                forward()
+            graph.instantiate()
+            graphs[name] = graph
+            status, _, count = cuda.cuGraphGetNodes(cuda.CUgraph(graph.raw_cuda_graph()))
+            assert status == cuda.CUresult.CUDA_SUCCESS
+            counts.append(count)
+        assert counts[0] - counts[1] == 1
+        bundle = PreparedGraphBundle(device=0, graphs=graphs, buffers={}, resources=(cases, compiled))
+        worker = bundle.start_worker()
+        try:
+            for step in range(2 * (width + 1)):
+                initial = torch.randint(0, 256, cases[0].storage.backing.shape, dtype=torch.uint8, device=device)
+                values = torch.randn_like(cases[0].values) * .25
+                values[0].zero_()
+                values[-1].fill_(2**-20)
+                lengths = [width if r == 0 else (step + r) % (width + 1) for r in range(requests)]
+                selected = [(step + r) % (requests + 2) + 1 for r in range(requests)]
+                starts = torch.tensor([127 + (step // (width + 1)) + r * 128
+                                       for r in range(requests)], device=device)
+                length_tensor = torch.tensor(lengths, dtype=torch.int32, device=device)
+                accepted = [min(step % (width + 1), length) for length in lengths]
+                tokens = torch.zeros((requests, width), dtype=torch.int64, device=device)
+                pages = torch.zeros((requests, 1), dtype=torch.int32, device=device)
+                for case in cases:
+                    case.storage.backing.copy_(initial)
+                    case.values.copy_(values)
+                    case.cache.fill_(173)
+                    case.counts.copy_(torch.tensor(accepted, dtype=torch.int32, device=device))
+                    case.metadata.set(starts, length_tensor, torch.tensor(selected, device=device),
+                                      pages, tokens, torch.tensor([requests], dtype=torch.int32, device=device))
+                torch.cuda.synchronize()
+                before = torch.cuda.memory_stats()
+                for name in ("separate", "fused") if step % 2 else ("fused", "separate"):
+                    ticket = worker.submit(bundle.graph_indices[name])
+                    completion = worker.poll(timeout_ms=120000)
+                    assert completion is not None and completion.ticket == ticket and completion.error is None
+                after = torch.cuda.memory_stats()
+                for key in ("allocation.all.allocated", "allocation.all.freed",
+                            "segment.all.allocated", "segment.all.freed"):
+                    assert before[key] == after[key]
+                control, candidate = cases
+                assert torch.equal(control.observed, candidate.observed)
+                assert torch.equal(control.storage.backing, candidate.storage.backing)
+                expected = initial.clone()
+                if mode == "verify":
+                    region = candidate.storage.layout.region("window")
+                    for request, count in enumerate(accepted):
+                        for token in range(count):
+                            destination = region.offset + ((int(starts[request]) + token) % 128) * 528
+                            source = (128 + token) * 528
+                            expected[selected[request], destination:destination + 528].copy_(
+                                candidate.observed[request, source:source + 528])
+                assert torch.equal(candidate.storage.backing, expected)
+                assert candidate.cache.eq(129).all() and candidate.values.eq(37).all()
+        finally:
+            worker.close()
+            for graph in graphs.values():
+                graph.reset()