Decode GLM width-one rounds in place without KDA records

docs/mtp-design.md

@@ -119,8 +119,12 @@ The shared native controller selects among the available counts. All ranks
 sample identical gathered distributions with identical device RNG state.
 History-copy proposals are not part of the GLM recipe.
 
-Target KDA layers retain one recurrent checkpoint per verification input.
-Commit selects the accepted checkpoint and convolution suffix. Indexed
+Rounds with drafts record each target KDA layer's per-token corrections without
+advancing the request's committed recurrent row. Commit replays the accepted
+records into that row and shifts the convolution history past the emitted
+tokens, so every round resumes from state column zero. Width one carries no
+drafts: its KDA layers decode the pending token in place through a one-column
+plan, and it performs no record replay or draft-tail restore. Indexed
 attention uses twelve-entry C4 tail rings so speculative writes preserve the
 preceding accepted tail. Prompt absorption and post-verification absorption
 write only the draft layer's latent KV and C4 state, using target feedback.

ignition/recipes/glm53/speculative.py

@@ -173,11 +173,14 @@ def __init__(self, owner, batch, width, name):
         self.sampling_scratch = prepare_sampling_scratch(b, w, o.requests.vocab_size, device, buffers=buffers)
         self.rows, self.columns = torch.arange(b, device=device), torch.arange(w, device=device)
         self.offsets = torch.arange(b * w, device=device)
-        self.state_columns = torch.arange(8, device=device)
+        # Width one carries no drafts to roll back, so its KDA layers decode in place through one state column.
+        recorded = w > 1
+        self.state_columns = torch.arange(8 if recorded else 1, device=device)
         self.convolution_regions = convolution_regions(o.pool)
         self.metadata = DecoderMetadata.allocate(b * w, b, blocks, device)
         self.metadata = DecoderMetadata(self.metadata.token_ids, torch.arange(b * w, device=device),
-            DeltaMetadata.allocate(b, device, verification_columns=8), self.metadata.indexed)
+            DeltaMetadata.allocate(b, device, verification_columns=self.state_columns.numel()),
+            self.metadata.indexed)
         self.validator = PreparedBatch(o.inputs, o.requests, DecoderMetadata.allocate(b, b, blocks, device),
             mode="decode", shard=o.weights.shard, pages=o.pool.memory_capacity[1:], shards=o.shards)
         self.validation = PreparedAllReduce(self.validator.validation_input, self.validator.validation_sum,
@@ -186,11 +189,13 @@ def __init__(self, owner, batch, width, name):
         self.decoder = PreparedDecoder(o.weights, o.scratch.rows(b * w, "decode", sample_rows=b * w),
             o.pool.decoder, self.metadata, mode="decode", session=o.session, name=f"{name}.target",
             workspace=o.workspace, max_cache_pages=o.max_cache_pages, context=o.context,
-            kda_records=o.kda_records, **o.collectives)
-        first = self.decoder.layers[o.kda_records.layers[0]].attention
-        self.kda_commit = PreparedKDACommit(o.pool, o.kda_records, first.recurrent_plan,
-            self.metadata.recurrent.state_indices[:, 0], self.metadata.recurrent.num_seqs,
-            session=o.session, workspace=o.workspace)
+            kda_records=o.kda_records if recorded else None, **o.collectives)
+        self.kda_commit = None
+        if recorded:
+            first = self.decoder.layers[o.kda_records.layers[0]].attention
+            self.kda_commit = PreparedKDACommit(o.pool, o.kda_records, first.recurrent_plan,
+                self.metadata.recurrent.state_indices[:, 0], self.metadata.recurrent.num_seqs,
+                session=o.session, workspace=o.workspace)
         self.draft_metadata = DraftMetadata(b, b, blocks, device, shards=o.shards)
         self.draft = o.draft(self.draft_metadata, f"{name}.draft") if w > 1 else None
         self.draft_logits = (buffers.zeros(self.draft.logits.shape, self.draft.logits.dtype)
@@ -206,7 +211,8 @@ def __init__(self, owner, batch, width, name):
         self.probabilities = buffers.zeros((b, w - 1, o.requests.vocab_size), torch.float32)
         self.valid = torch.zeros((b, w - 1), dtype=torch.bool, device=device)
         self.seed, self.offset, self.base_offset, self.sampled = (torch.zeros_like(self.before) for _ in range(4))
-        self.tail_backup = torch.empty_like(o.pool.draft.index.tail[:b])
+        # Draft passes advance the draft C4 tail ring, which is restored before absorbing the accepted tokens.
+        self.tail_backup = torch.empty_like(o.pool.draft.index.tail[:b]) if self.draft is not None else None
         self.broadcast_buffers, self.name = buffers, name
         self.operations = None
 
@@ -256,8 +262,9 @@ def target_metadata(self, slots, limits):
         g.convolution_slots.copy_(torch.where(self.active, states[:, 0], -1))
         g.fresh.zero_()
         g.accepted_tokens.fill_(1)
-        # Verification reads each request's live row in column zero and records its tokens.
-        states = torch.where(self.state_columns < 1, states, 0)
+        # Verification reads each request's live row in column zero and records its tokens; width one
+        # advances that row in place.
+        states = torch.where(self.state_columns < 1, states[:, :self.state_columns.numel()], 0)
         g.state_indices.copy_(torch.where(self.active[:, None], states, 0))
         g.num_seqs.copy_(o.inputs.live_rows.clamp(0, b))
         g.num_tokens.copy_(ends[-1:])
@@ -276,20 +283,25 @@ def target_metadata(self, slots, limits):
         m.sample_indices.copy_((starts[:, None] + self.columns).clamp(0, b * w - 1).reshape(-1))
 
     def commit_state(self, slots, counts):
-        """Replay each live KDA row's accepted records for every layer, then shift its convolution history."""
+        """Replay each live KDA row's accepted records for every layer, then shift its convolution history.
+
+        Width one has already advanced each live KDA row, and its single token leaves the convolution
+        history in place, so it saves only the target output.
+        """
         o, w = self.owner, self.width
         states = o.inputs.request_state_rows[slots.clamp(0, o.requests.capacity - 1)]
         live = torch.where(counts > 0, states[:, 0], 0).long()
-        self.kda_commit(counts)
-        backing = o.pool.backing
-        halves, channels, columns = self.convolution_regions
-        args = (backing.view(torch.bfloat16), slots, counts, o.inputs.request_state_rows, halves,
-                backing.shape[1] // 2, channels, columns, 128)
-        grid = (halves.numel(), self.batch, triton.cdiv(channels, 128))
-        if torch.compiler.is_compiling() or get_proxy_mode() is not None:
-            torch.library.wrap_triton(_shift_convolution)[grid](*args, num_warps=4)
-        else:
-            _shift_convolution[grid](*args, num_warps=4)
+        if self.kda_commit is not None:
+            self.kda_commit(counts)
+            backing = o.pool.backing
+            halves, channels, columns = self.convolution_regions
+            args = (backing.view(torch.bfloat16), slots, counts, o.inputs.request_state_rows, halves,
+                    backing.shape[1] // 2, channels, columns, 128)
+            grid = (halves.numel(), self.batch, triton.cdiv(channels, 128))
+            if torch.compiler.is_compiling() or get_proxy_mode() is not None:
+                torch.library.wrap_triton(_shift_convolution)[grid](*args, num_warps=4)
+            else:
+                _shift_convolution[grid](*args, num_warps=4)
         hidden = self.decoder.scratch.hidden.view(self.batch, w, 4096)
         selected_hidden = hidden[self.rows, (counts - 1).clamp_min(0).long()]
         save_hidden(selected_hidden, o.pool.last_hidden,
@@ -315,7 +327,8 @@ def __call__(self):
         self.before.copy_(o.requests.consumed_tokens[safe])
         self.pending.copy_(o.requests.pending_token[safe])
         gather_rows(o.pool.last_hidden, o.inputs.request_state_rows[safe, 0], self.previous)
-        gather_rows(o.pool.draft.index.tail, o.inputs.request_state_rows[safe, 0], self.tail_backup)
+        if self.tail_backup is not None:
+            gather_rows(o.pool.draft.index.tail, o.inputs.request_state_rows[safe, 0], self.tail_backup)
         self.base_offset.copy_(torch.where(o.draft_generation[safe] == o.requests.generation[safe], o.draft_offset[safe], 0))
         self.seed.copy_(o.requests.seed.view(torch.int64)[safe] ^ DRAFT_RNG_DOMAIN)
         for i, proposal in enumerate(proposals):
@@ -346,9 +359,10 @@ def __call__(self):
         verify()
         commit()
         self.commit_state(slots, commit.emitted_count)
-        tail = o.pool.draft.index.tail
-        save_hidden(self.tail_backup.view(b, -1), tail.view(tail.shape[0], -1),
-                    o.inputs.request_state_rows[safe, 0], self.active)
+        if self.tail_backup is not None:
+            tail = o.pool.draft.index.tail
+            save_hidden(self.tail_backup.view(b, -1), tail.view(tail.shape[0], -1),
+                        o.inputs.request_state_rows[safe, 0], self.active)
         hidden = torch.cat((self.previous[:, None], self.decoder.scratch.hidden.view(b, w, 4096)[:, :-1]), dim=1)
         self.absorb_metadata.set(slots, commit.emitted_count.long(), self.before - 1,
             self.candidates, hidden, o.draft_tables, state_rows=o.inputs.request_state_rows, capacity=w)

tests/recipes/test_glm53_mtp_state_cuda.py

@@ -2,7 +2,10 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 Flare community contributors
 # SPDX-License-Identifier: Apache-2.0
 
-"""Accepted KDA prefixes, replayed from verification records, match plain decode across widths."""
+"""Accepted KDA prefixes, replayed from verification records, match plain decode across widths.
+
+Width-one rounds advance the live rows in place, and recorded rounds resume from their results.
+"""
 
 import os
 from pathlib import Path
@@ -26,15 +29,34 @@
                                reason="requires an explicitly assigned CUDA device")
 
 
-@torch.no_grad()
-def test_learned_kda_verification_commits_each_prefix_before_plain_decode():
+def first_kda_layer():
     snapshot = Path(os.environ.get("FLARE_GLM53_SNAPSHOT",
         "/data/checkpoint-releases/GLM-5.3-Flash-NVFP4-MXFP8-CSF-QAD"))
     manifest = inspect_indexed_checkpoint(snapshot)
     module = create_model_tree(snapshot).model.language_model.layers[0].self_attn
     with GlmWeightLoader(manifest, "cuda:0", shard=TensorParallelShard(0, 4),
                          prefixes=(f"{TEXT_PREFIX}.layers.0.self_attn.",)) as loader:
-        weights = loader.load_attention(0)
+        return module, loader.load_attention(0)
+
+
+def round_state(pool, width, rows, kda_commit=None):
+    """A round's state commit over `rows`: requests in slots 0, 1, ... own their rows' column zero."""
+    op = PreparedRound.__new__(PreparedRound)
+    batch = rows.shape[0]
+    op.owner = SimpleNamespace(pool=pool, requests=SimpleNamespace(capacity=batch),
+        inputs=SimpleNamespace(row_slot_ids=torch.arange(batch, dtype=torch.int32, device="cuda:0"),
+                               request_state_rows=rows))
+    op.width, op.batch, op.rows = width, batch, torch.arange(batch, device="cuda:0")
+    op.convolution_regions = convolution_regions(pool)
+    op.kda_commit = kda_commit
+    op.decoder = SimpleNamespace(scratch=SimpleNamespace(
+        hidden=torch.zeros((batch * width, 4096), dtype=torch.bfloat16, device="cuda:0")))
+    return op
+
+
+@torch.no_grad()
+def test_learned_kda_verification_commits_each_prefix_before_plain_decode():
+    module, weights = first_kda_layer()
     pool = ContinuationPool.allocate(("linear_attention",), 4, 10, 1, "cuda:0", mtp=True)
     state = pool.decoder.layers[0]
     arena = ScratchArena("cuda:0")
@@ -69,6 +91,8 @@ def test_learned_kda_verification_commits_each_prefix_before_plain_decode():
             meta.state_indices.copy_(torch.tensor([[1, 0, 0, 0, 0, 0, 0, 0]]))
             meta.num_seqs.fill_(1)
             meta.num_tokens.fill_(width)
+        # Plain decode first runs after the session freezes, so it resolves its kernels here.
+        operators[1](sources[1])
         generator = torch.Generator(device="cuda:0").manual_seed(531)
         chain = torch.randn((7, 4096), generator=generator, device="cuda:0", dtype=torch.bfloat16) * .1
         state.recurrent.normal_(0, .01, generator=generator)
@@ -118,3 +142,76 @@ def test_learned_kda_verification_commits_each_prefix_before_plain_decode():
             torch.testing.assert_close(state.convolution[1, :, :3], expected_c, rtol=.025, atol=.002)
             actual_output = operators[1](sources[1])
             torch.testing.assert_close(actual_output, expected_output, rtol=.025, atol=.005)
+
+
+@torch.no_grad()
+def test_width_one_advances_live_rows_in_place_between_recorded_rounds():
+    module, weights = first_kda_layer()
+    pool = ContinuationPool.allocate(("linear_attention",), 4, 10, 1, "cuda:0", mtp=True)
+    state = pool.decoder.layers[0]
+    arena = ScratchArena("cuda:0")
+    # Request zero owns row 1. Request one's row 2 stays inactive, as a rejected or idle request does.
+    rows = torch.tensor([[1, 0, 0, 0, 0, 0, 0, 0], [2, 0, 0, 0, 0, 0, 0, 0]], device="cuda:0")
+    # Width one decodes in place through one state column. Its recorded form, with the same projection
+    # capacity, is the reference: recorded rounds verify against eight columns and commit accepted records.
+    shapes = {"in_place": (1, 1), "recorded": (1, 8), "verification": (6, 8)}
+    metas = {key: DeltaMetadata.allocate(2, "cuda:0", verification_columns=columns)
+             for key, (_, columns) in shapes.items()}
+    sources = {key: torch.zeros((2 * width, 4096), dtype=torch.bfloat16, device="cuda:0")
+               for key, (width, _) in shapes.items()}
+    records = KDARecords.allocate({0: weights}, state.recurrent.shape[0], state.recurrent.shape[1], "cuda:0")
+    with RecordedPreparationSession(device=torch.device("cuda:0"), checkpoint="glm53-mtp-state",
+            workload_schema="glm-kda-width-one-v1", numerical_policy="default") as session:
+        session.register_scratch(arena)
+        operators = {key: PreparedGlmKDA(module, weights, sources[key],
+            KDAScratch.allocate(2 * width, 16, "cuda:0"), state, metas[key],
+            mode="decode", session=session, name=key, workspace=arena,
+            records=None if key == "in_place" else records.layer(0)) for key, (width, _) in shapes.items()}
+        assert operators["in_place"].records is None and operators["in_place"].columns == 1
+        rounds = {key: round_state(pool, width, rows, None if key == "in_place" else
+            PreparedKDACommit(pool, records, operators[key].recurrent_plan, metas[key].state_indices[:, 0],
+                              metas[key].num_seqs, session=session, workspace=arena))
+            for key, (width, _) in shapes.items()}
+        session.tune()
+        counts = torch.zeros(2, dtype=torch.int32, device="cuda:0")
+
+        def run(key, tokens, accepted):
+            # Metadata as PreparedRound.target_metadata packs it with request one inactive.
+            meta, length = metas[key], tokens.shape[0]
+            meta.query_start_loc.copy_(torch.tensor([0, length, length], dtype=torch.int32))
+            meta.convolution_slots.copy_(torch.tensor([1, -1]))
+            meta.state_indices.zero_()
+            meta.state_indices[0, 0] = 1
+            meta.num_seqs.fill_(2)
+            meta.num_tokens.fill_(length)
+            sources[key].zero_()
+            sources[key][:length].copy_(tokens)
+            output = operators[key](sources[key])[:length].clone()
+            counts.copy_(torch.tensor([accepted, 0], dtype=torch.int32))
+            rounds[key].commit_state(rounds[key].owner.inputs.row_slot_ids.long(), counts)
+            return output
+
+        generator = torch.Generator(device="cuda:0").manual_seed(977)
+        chain = torch.randn((7, 4096), generator=generator, device="cuda:0", dtype=torch.bfloat16) * .1
+        state.recurrent.normal_(0, .01, generator=generator)
+        state.convolution.normal_(0, .1, generator=generator)
+        initial_r, initial_c = state.recurrent.clone(), state.convolution.clone()
+        # Each sequence runs width one, a recorded round that accepts three of six candidates, and width one.
+        results = {}
+        for key in ("recorded", "in_place"):
+            state.recurrent.copy_(initial_r)
+            state.convolution.copy_(initial_c)
+            outputs = (run(key, chain[:1], 1), run("verification", chain[1:7], 3)[:3], run(key, chain[4:5], 1))
+            results[key] = outputs, state.recurrent.clone(), state.convolution.clone()
+            # The inactive request's row and the null row are unchanged by every round.
+            for row in (0, 2):
+                torch.testing.assert_close(state.recurrent[row], initial_r[row], rtol=0, atol=0)
+                torch.testing.assert_close(state.convolution[row], initial_c[row], rtol=0, atol=0)
+        (expected, expected_r, expected_c), (actual, actual_r, actual_c) = results["recorded"], results["in_place"]
+        # Convolution does not depend on recording; recording rounds the recurrence to BF16 before the gated
+        # norm, as b12x's recording does, and replays the committed state from its records.
+        torch.testing.assert_close(actual_c, expected_c, rtol=0, atol=0)
+        torch.testing.assert_close(actual_r[1], expected_r[1], rtol=.025, atol=.002)
+        for value, reference in zip(actual, expected, strict=True):
+            relative = ((value.float() - reference.float()).norm() / reference.float().norm()).item()
+            assert relative < 1e-2, relative

tests/recipes/test_glm53_speculative.py

@@ -44,8 +44,8 @@ def test_shared_commit_without_previous_token_embeddings():
 
 
 @pytest.mark.parametrize("width,limits,positions,token_slots,state_indices", [
-    (1, [1, 1], [255, 510], [10 * 256 + 255, 21 * 256 + 254],
-     [[1, 0, 0, 0, 0, 0, 0, 0], [2, 0, 0, 0, 0, 0, 0, 0]]),
+    # Width one decodes each live row in place through a single state column.
+    (1, [1, 1], [255, 510], [10 * 256 + 255, 21 * 256 + 254], [[1], [2]]),
     (6, [6, 2], [255, 256, 257, 258, 259, 260, 510, 511, -1, -1, -1, -1],
      [10 * 256 + 255, *range(11 * 256, 11 * 256 + 5), 21 * 256 + 254, 21 * 256 + 255],
      # Verification records tokens against each request's live row only.
@@ -61,11 +61,12 @@ def test_verification_packs_unequal_lengths_and_retains_request_output_rows(
     inputs.request_state_rows.copy_(torch.tensor([[1, 8, 9, 10, 11, 12, 13, 14], [2, 15, 16, 17, 18, 19, 20, 21]]))
     op.owner = SimpleNamespace(inputs=inputs, requests=SimpleNamespace(capacity=2), shards=ContextShards())
     m = DecoderMetadata.allocate(2 * width, 2, 2, "cpu")
+    columns = len(state_indices[0])
     op.metadata = DecoderMetadata(m.token_ids, torch.zeros(2 * width, dtype=torch.int64),
-        DeltaMetadata.allocate(2, "cpu", verification_columns=8), m.indexed)
+        DeltaMetadata.allocate(2, "cpu", verification_columns=columns), m.indexed)
     op.active = torch.tensor([True, True])
     op.before = torch.tensor([255, 510])
-    op.offsets, op.columns, op.state_columns = torch.arange(2 * width), torch.arange(width), torch.arange(8)
+    op.offsets, op.columns, op.state_columns = torch.arange(2 * width), torch.arange(width), torch.arange(columns)
     op.candidates = torch.arange(2 * width).view(2, width)
     op.target_metadata(torch.tensor([0, 1]), torch.tensor(limits, dtype=torch.int32))
     assert op.metadata.recurrent.query_start_loc.tolist() == [0, limits[0], sum(limits)]