From 8aab9c9e71adbb3e12788d6d12376c5f4af29daf Mon Sep 17 00:00:00 2001 From: Jason Cook Date: Thu, 5 Nov 2026 13:12:00 -0500 Subject: [PATCH] Decode GLM width-one rounds in place without KDA records A GLM width-one speculative round verifies only the pending target token, so it has no drafts to roll back. Its target KDA layers recorded per-token corrections, and a record replay then committed one record into each live row; the round also saved and restored the draft C4 tail ring, which no draft pass touches at width one. Width-one rounds now prepare their KDA layers with the ordinary one-column decode plan, which advances each live recurrent row in place. Inactive rows keep the null state index and convolution slot -1, so their recurrent state and convolution history do not change. Width-one rounds construct no record replay and skip the convolution shift, which leaves history unchanged when a round emits at most one token, and they skip the draft tail backup and restore. They still absorb the emitted token into the draft cache and save the target output, so the following round can draft. Wider rounds are unchanged. docs/mtp-design.md describes the GLM record and in-place KDA paths, replacing a description of per-input recurrent checkpoints that the recipe does not use. Compatibility: no package metadata or prepared-contract change. Exported packages keep their captured graphs; re-export prepares the width-one graphs. A request whose width-one target distribution fails validation is marked as an error and emits no token. Its KDA row now advances by the verified token, as its convolution history already did; Helios publishes no prefix for requests with an error code. Validation on main (972e28e4) with this change, on one RTX PRO 6000 (SM120), with layer 0 KDA weights of GLM-5.3-Flash-NVFP4-MXFP8-CSF-QAD at the TP4 head partition: - tests/recipes/test_glm53_mtp_state_cuda.py: 2 passed. A sequence of a width-one round, a recorded six-candidate round accepting three tokens, and a width-one round matches the same sequence with the recorded width-one plan: convolution history is bitwise equal, the recurrent row is within rtol 0.025 and atol 0.002, and the round outputs are within 1e-2 relative norm. The inactive request's row and the null row stay bitwise unchanged. test_learned_kda_verification_commits_each_prefix_before_plain_decode now resolves its plain-decode kernels before the preparation session freezes; without that step it fails on main with KernelResolutionFrozenError. - tests/recipes/test_glm53_speculative.py with the other GLM recipe, speculative-capacity and prepared-contract CPU tests: 222 passed, 47 skipped. Six tests/recipes/test_declarations.py cases fail because the Kimi and Qwen checkpoints they read are absent on the test host; they fail identically without this change. Width-one target metadata packs one state column. - DGX Spark (GB10, SM121): GLM-5.3-Flash packages exported from a tree that contains this change and other GLM recipe changes, without context parallelism, served at TP2 on two hosts and TP4 on four hosts with five drafts and adaptive verification. Each package passed a 32-prompt record-retrieval and counting check serially and at eight concurrent requests, with identical answers in both modes. Not run: a complete width-one round graph compared against ordinary decode, SM103 hardware, and a decode-time measurement that isolates this change. Design-Signed-Off-By: flare-02; design-review=4d3a --- docs/mtp-design.md | 8 +- ignition/recipes/glm53/speculative.py | 64 ++++++++----- tests/recipes/test_glm53_mtp_state_cuda.py | 105 ++++++++++++++++++++- tests/recipes/test_glm53_speculative.py | 9 +- 4 files changed, 151 insertions(+), 35 deletions(-) diff --git a/docs/mtp-design.md b/docs/mtp-design.md index 90c7c9aa..dd991add 100644 --- a/docs/mtp-design.md +++ b/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. diff --git a/ignition/recipes/glm53/speculative.py b/ignition/recipes/glm53/speculative.py index 1540a0c1..8b08e04b 100644 --- a/ignition/recipes/glm53/speculative.py +++ b/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) diff --git a/tests/recipes/test_glm53_mtp_state_cuda.py b/tests/recipes/test_glm53_mtp_state_cuda.py index e7d47acd..c20994b7 100644 --- a/tests/recipes/test_glm53_mtp_state_cuda.py +++ b/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 diff --git a/tests/recipes/test_glm53_speculative.py b/tests/recipes/test_glm53_speculative.py index 58ca0c38..dc1f12d4 100644 --- a/tests/recipes/test_glm53_speculative.py +++ b/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)]