Skip to content

Commit 5cee115

Browse files
authored
Reuse dead GLM prefill storage to enlarge the KV cache (#253)
On a DGX Spark, GLM-5.3 TP2 serving sizes its KV cache from the memory left after the core's fixed mappings, so every MiB of retained scratch is cache capacity. Three GLM preparation buffers duplicate storage whose lifetimes never overlap: - Each serialized prefill variant allocated its own BF16 [capacity, 4096] MTP target-feedback buffer, although prefill variants never run concurrently. - Each prefill variant's draft metadata allocated a BF16 [capacity, 4096] hidden input, although the target generation's attention_hc.normalized buffer of the same shape is dead once target generation returns and before draft absorption reads the hidden input. - Context-parallel MLA prefill allocated separate query and expanded-value buffers, although each 2048-row slice consumes its query rows in absorption before value expansion writes the same rows. prepare_rank now allocates one feedback buffer for the largest prefill capacity and binds a row view of it per variant. DraftMetadata accepts caller-owned hidden storage, validating its shape, dtype, device and contiguity, and PreparedPrefill passes the target's attention_hc.normalized buffer. MLAScratch aliases expanded values to the query buffer in context-parallel prefill only; every other mode keeps separate storage. Operations, shapes, inputs and execution order are unchanged, so outputs are bit-identical. The result is a smaller core and a larger KV cache at the same settings. For GLM-5.3-Flash FP4-Local TP2 with context parallelism, concurrency 8, MTP and 4096-token prefill chunks on two DGX Sparks, the export inventory reports 88,468 MiB of core fixed mappings on rank 0, against 88,758 MiB without this change (-290 MiB: -179 MiB disposable scratch, -115 MiB other owned tensors), and serving allocates 4,456,448 KV tokens, against 4,384,768 (+71,680, +1.6%). Evidence: test_glm53_mla_cuda compares compiled context-parallel slice execution with aliased and separate value storage over poisoned buffers, alternating 8192- and 2049-row prefills; test_glm53_speculative and test_glm53_draft_metadata_cuda cover poisoned caller-owned hidden storage and invalid storage rejection; test_glm53_stack covers the alias and the capacity views. On a DGX Spark at main bbad44f, the GLM recipe suites (stack, speculative, draft metadata, MLA, MTP state) pass 42 tests with this change; test_learned_kda_verification_commits_each_prefix_before_plain_decode fails with and without it. Serving cores exported with and without this change produce identical first-token log-probabilities for fixed 8K and 128K prompts. Cross-host replay of the 8192-row chunk configuration and outputs beyond the first token were not compared: sampled generations are not reproducible across requests even within one server, and 500K prefill log-probabilities vary between runs of the same core. Design-Signed-Off-By: yatesdr; design-review=cfcf
1 parent 6bd20bc commit 5cee115

8 files changed

Lines changed: 142 additions & 8 deletions

File tree

‎ignition/recipes/glm53/bootstrap.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -482,6 +482,8 @@ def add(name, function, kind, rows, tokens, fields, *, enabled=None):
482482
forward = PreparedVisionGeneration(encoder, vision_inputs, vision_state, completion, frame, sequence)
483483
add(name, forward, "vision", 1, capacity, (*vision_inputs.buffers(), "vision_feature_slot"))
484484
vision_owners.extend((vision_metadata, encoder, forward))
485+
prefill_feedback = (torch.empty((config.prefill_chunk_size, 4096), dtype=torch.bfloat16, device=device)
486+
if config.mtp else None)
485487
for enabled in graph_modes:
486488
suffix = "-no-mtp" if config.mtp and not enabled else ""
487489
mode_inputs = tuple(field for field in forward_inputs if enabled or field != "draft_page_tables")
@@ -493,7 +495,7 @@ def add(name, function, kind, rows, tokens, fields, *, enabled=None):
493495
metadata[name] = DecoderMetadata.allocate(capacity, rows, blocks, device)
494496
storage = prefill_storage if mode == "prefill" else decoder_storage
495497
scratch[name] = storage.rows(capacity, mode, sample_rows=rows)
496-
feedback = (torch.empty((capacity, 4096), dtype=torch.bfloat16, device=device)
498+
feedback = (prefill_feedback[:capacity]
497499
if enabled and kind != "scored_prefill" else None)
498500
decoders[name] = PreparedDecoder(weights, scratch[name], pool.decoder, metadata[name], mode=mode,
499501
scoring=scoring if kind == "scored_prefill" else None,

‎ignition/recipes/glm53/mla.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,9 @@ def buffer(*shape):
7878
# Fused layer-input projection: query 1536, latent 512, index key 128 and index head weights 32.
7979
projected = buffer(2208)
8080
prefill = context_ranks > 1 and mode == "prefill"
81+
query = buffer(heads, 256)
82+
# CP prefill consumes each query slice before expanding values into the same rows.
83+
expanded = query if prefill else buffer(heads, 256)
8184
exchange = min(tokens, SLICE_ROWS) if context_ranks > 1 and not prefill else 0
8285
# Context-parallel rows absorb, attend and expand one slice at a time; otherwise every row at once.
8386
absorbed = exchange or (min(tokens, PREFILL_SLICE_ROWS) if prefill else tokens)
@@ -90,7 +93,7 @@ def buffer(*shape):
9093

9194
window = min(WINDOW_RECORDS, blocks * 256) if remote and prefill else 0
9295
return cls(projected, projected[:, :1536], buffer(1536), projected[:, 1536:2048], buffer(512),
93-
buffer(heads, 256), zeros(absorbed, heads, 512), zeros(absorbed, heads, 512), buffer(heads, 256),
96+
query, zeros(absorbed, heads, 512), zeros(absorbed, heads, 512), expanded,
9497
buffer(4096), zeros(tokens, dtype=torch.int32), index,
9598
zeros(exchange, heads * context_ranks, 512), zeros(exchange, heads * context_ranks, 512),
9699
zeros(exchange, heads * context_ranks, dtype=torch.float32), zeros(positions, 528, dtype=torch.uint8),

‎ignition/recipes/glm53/mtp.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -126,13 +126,18 @@ def _draft_metadata(Lengths, Starts, Slots, Tokens, Hidden, Tables, StateRows, R
126126

127127

128128
class DraftMetadata:
129-
def __init__(self, tokens, rows, blocks, device, *, shards=None):
129+
def __init__(self, tokens, rows, blocks, device, *, shards=None, hidden=None):
130130
from .dcp import ContextShards
131131

132132
self.shards = ContextShards() if shards is None else shards
133133
self.indexed = IndexMetadata.allocate(tokens, rows, blocks, device)
134134
self.token_ids = torch.zeros(tokens, dtype=torch.int64, device=device)
135-
self.hidden = torch.zeros((tokens, 4096), dtype=torch.bfloat16, device=device)
135+
if hidden is None:
136+
hidden = torch.zeros((tokens, 4096), dtype=torch.bfloat16, device=device)
137+
if (hidden.shape != (tokens, 4096) or hidden.dtype != torch.bfloat16
138+
or hidden.device != self.token_ids.device or not hidden.is_contiguous()):
139+
raise ValueError("GLM draft hidden storage must be contiguous BF16 [tokens, 4096] on the metadata device")
140+
self.hidden = hidden
136141
self.indices = torch.arange(tokens, device=device)
137142
self.output_indices = torch.zeros(rows, dtype=torch.int64, device=device)
138143
self.num_tokens = torch.zeros(1, dtype=torch.int32, device=device)

‎ignition/recipes/glm53/speculative.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,9 @@ def __init__(self, owner, generation, feedback, name):
133133
self.owner, self.generation, self.feedback = owner, generation, feedback
134134
tokens = feedback.shape[0]
135135
device = feedback.device
136-
self.metadata = DraftMetadata(tokens, 1, owner.inputs.page_tables.shape[1], device, shards=owner.shards)
136+
# Target generation finishes before draft metadata overwrites its normalized attention input.
137+
self.metadata = DraftMetadata(tokens, 1, owner.inputs.page_tables.shape[1], device, shards=owner.shards,
138+
hidden=generation.decoder.scratch.attention_hc.normalized)
137139
self.absorb = owner.draft(self.metadata, f"{name}.absorb", cache_only=True, mode="prefill")
138140
self.previous = torch.zeros((1, 4096), dtype=torch.bfloat16, device=device)
139141
self.indices = torch.arange(tokens, device=device)

‎tests/recipes/test_glm53_draft_metadata_cuda.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,8 @@ def test_fused_draft_metadata_matches_reference(world, capacity, requests):
2323
generator = torch.Generator().manual_seed(world * 10 + capacity)
2424
blocks, slots_capacity = 40, 8
2525
shards = ContextShards.on(1, world, device) if world > 1 else ContextShards()
26-
fused = DraftMetadata(requests * capacity, requests, blocks, device, shards=shards)
26+
storage = torch.full((requests * capacity, 4096), float("nan"), dtype=torch.bfloat16, device=device)
27+
fused = DraftMetadata(requests * capacity, requests, blocks, device, shards=shards, hidden=storage)
2728
reference = DraftMetadata(requests * capacity, requests, blocks, device, shards=shards)
2829
slots = torch.tensor([3, 0, 5, 6][:requests], dtype=torch.int64, device=device)
2930
lengths = torch.randint(0, capacity + 1, (requests,), generator=generator).to(device)

‎tests/recipes/test_glm53_mla_cuda.py‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,14 +6,18 @@
66

77
import os
88
import time
9+
from dataclasses import replace
910
from pathlib import Path
11+
from types import SimpleNamespace
1012

1113
import pytest
1214
import torch
1315
from b12x.attention import sparse_mla
1416
from cuda.bindings import driver as cuda
1517

1618
from ignition.components.parallel import TensorParallelShard
19+
from ignition.components.linear import BatchedQuantizedProjection, BatchedUnquantizedProjection
20+
from ignition.components.sparse_mla import PreparedSparseMLA
1721
from ignition.preparation.checkpoint import inspect_indexed_checkpoint
1822
from ignition.preparation.compilation import compile_forward
1923
from ignition.preparation.execution import GraphSignature, PreparedExecution
@@ -235,3 +239,86 @@ def forward():
235239
if worker is not None:
236240
worker.close()
237241
graph.reset()
242+
243+
244+
@pytest.mark.parametrize("model_weights", [2], indirect=True)
245+
@torch.no_grad()
246+
def test_cp_prefill_query_reuse_matches_separate_values_across_slices(model_weights):
247+
from b12x.attention._shared.mla.reference import pack_mla_kv_cache_reference
248+
from ignition.recipes.glm53.dcp import ContextShards, PREFILL_SLICE_ROWS
249+
250+
module, weights, world = model_weights
251+
device, capacity, blocks, heads = torch.device("cuda:0"), 8192, 32, 64 // world
252+
torch.manual_seed(53)
253+
caches = [pack_mla_kv_cache_reference(torch.randn(blocks * 256, 512, device=device)).view(-1, 64, 528)
254+
for _ in range(world)]
255+
peers = torch.tensor([cache.data_ptr() for cache in caches], dtype=torch.int64, device=device)
256+
storage = MLAScratch.allocate(capacity, heads, blocks, "prefill", device, context_ranks=world)
257+
queries = torch.randn(storage.query.shape, dtype=torch.bfloat16, device=device) * .1
258+
arena, operations, graphs, compiled = ScratchArena(device), {}, {}, {}
259+
kv = weights["kv_b_proj"]
260+
projection = BatchedUnquantizedProjection if isinstance(kv, torch.Tensor) else BatchedQuantizedProjection
261+
if isinstance(kv, torch.Tensor):
262+
values = kv.view(heads, 512, 512)
263+
absorb_weights, expand_weights = values[:, :256], values[:, 256:]
264+
else:
265+
values, scales = kv.weight.values.view(heads, 512, 512), kv.weight.scale_rows.view(heads, 512, 16)
266+
absorb_weights, expand_weights = (values[:, :256], scales[:, :256]), (values[:, 256:], scales[:, 256:])
267+
with RecordedPreparationSession(device=device, checkpoint="glm53-cp-scratch",
268+
workload_schema="glm-cp-scratch-v1", numerical_policy="default") as session:
269+
session.register_scratch(arena)
270+
for rows in (capacity, 2049):
271+
for alias in (False, True):
272+
key = rows, alias
273+
s = storage.rows(rows, "prefill")
274+
if not alias:
275+
s = replace(s, expanded=torch.empty_like(s.expanded))
276+
metadata = IndexMetadata.allocate(rows, 1, blocks, device)
277+
metadata.positions.copy_(torch.arange(rows, device=device) + 8)
278+
metadata.block_table[0].copy_(torch.arange(blocks, device=device))
279+
s.index.physical.fill_(-1)
280+
s.index.physical[:, :4].copy_(torch.arange(4, device=device))
281+
s.index.selected_lengths.fill_(4)
282+
# Isolate the production CP slice loop with real projections, attention and peer-cache reads.
283+
op = PreparedGlmMLA.__new__(PreparedGlmMLA)
284+
op.scratch, op.metadata, op.state = s, metadata, SimpleNamespace(peers=peers)
285+
op.context = SimpleNamespace(shards=ContextShards.on(0, world, device))
286+
op.windows, op.slices = (), []
287+
caps = sparse_mla.Caps(device=device, num_q_heads=heads, max_q_rows=PREFILL_SLICE_ROWS, max_width=2051,
288+
softmax_scale=module.scaling, max_batch=PREFILL_SLICE_ROWS, max_kv_rows=s.slab.shape[0],
289+
kv_dtype=torch.uint8, head_dim=512, v_head_dim=512, page_size=64,
290+
model_type=sparse_mla.ModelType.GLM_NEXT, cache_record_bytes=528, mode="extend")
291+
for start in range(0, rows, PREFILL_SLICE_ROWS):
292+
end = min(start + PREFILL_SLICE_ROWS, rows)
293+
absorbed, attended = s.absorbed_query[:end - start], s.attended[:end - start]
294+
name = f"cp{rows}.{alias}.{start}"
295+
absorb = projection(s.query[start:end].transpose(0, 1), absorb_weights,
296+
absorbed.transpose(0, 1), major="n", session=session, name=name + ".absorb")
297+
attention = PreparedSparseMLA(absorbed, s.slab.view(-1, 64, 528), s.index.physical[start:end],
298+
s.visible[start:end], s.index.selected_lengths[start:end], attended, caps=caps,
299+
session=session, name=name + ".attention", workspace=arena)
300+
expand = projection(attended.transpose(0, 1), expand_weights,
301+
s.expanded[start:end].transpose(0, 1), major="k", session=session, name=name + ".expand")
302+
op.slices.append((start, end, absorb, attention, expand))
303+
operations[key] = op
304+
session.tune()
305+
for key, op in operations.items():
306+
def forward(op=op):
307+
op.scratch.query.copy_(queries[:op.scratch.query.shape[0]])
308+
op.attend_gathered_context()
309+
arena.tensor.fill_(255)
310+
compiled[key] = compile_forward(forward)
311+
compiled[key]()
312+
graphs[key] = torch.cuda.CUDAGraph()
313+
with session.capture(), torch.cuda.graph(graphs[key]):
314+
compiled[key]()
315+
session.freeze()
316+
for rows in (capacity, 2049, capacity):
317+
graphs[rows, False].replay()
318+
expected = operations[rows, False].scratch.expanded.clone()
319+
storage.query.fill_(float("nan"))
320+
graphs[rows, True].replay()
321+
torch.cuda.synchronize(device)
322+
assert torch.equal(operations[rows, True].scratch.expanded, expected)
323+
for graph in graphs.values():
324+
graph.reset()

‎tests/recipes/test_glm53_speculative.py‎

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717
from ignition.recipes.glm53.dcp import ContextShards
1818
from ignition.recipes.glm53.speculative import PreparedRound
1919
from ignition.recipes.glm53.mtp import DraftMetadata
20+
from ignition.recipes.glm53.stack import DecoderScratch
21+
from ignition.components.parallel import TensorParallelShard
2022
from ignition.recipes.glm53.state import DecoderMetadata
2123
from ignition.components.recurrent import DeltaMetadata
2224

@@ -76,8 +78,12 @@ def test_verification_packs_unequal_lengths_and_retains_request_output_rows(
7678
assert op.metadata.recurrent.state_indices.tolist() == state_indices
7779

7880

79-
def test_draft_rows_exclude_exhausted_requests_and_restore_output_order():
80-
metadata = DraftMetadata(3, 3, 2, "cpu")
81+
@pytest.mark.parametrize("reuse", [False, True])
82+
def test_draft_rows_exclude_exhausted_requests_and_restore_output_order(reuse):
83+
scratch = DecoderScratch.allocate(3, 2, "prefill", "cpu", shard=TensorParallelShard(0, 2), sample_rows=1)
84+
storage = scratch.attention_hc.normalized if reuse else None
85+
metadata = DraftMetadata(3, 3, 2, "cpu", hidden=storage)
86+
metadata.hidden.fill_(float("nan"))
8187
slots = torch.tensor([2, 0, 1])
8288
state_rows = torch.arange(1, 4)[:, None].expand(3, 8)
8389
tables = torch.tensor([[10, 11], [20, 21], [30, 31]])
@@ -89,13 +95,25 @@ def test_draft_rows_exclude_exhausted_requests_and_restore_output_order():
8995
assert metadata.indexed.token_slots.tolist() == [10 * 256 + 255, 21 * 256 + 254, -1]
9096
assert metadata.token_ids[:2].tolist() == [8, 9]
9197
assert metadata.hidden[metadata.output_indices][1:, 0].tolist() == [1, 2]
98+
assert hidden[:, 0].tolist() == [0, 1, 2]
99+
if reuse:
100+
assert torch.equal(storage, metadata.hidden)
101+
assert storage.data_ptr() == metadata.hidden.data_ptr()
92102
metadata.set(slots, torch.zeros(3, dtype=torch.int64), torch.full((3,), 512),
93103
torch.tensor([7, 8, 9]), hidden, tables, capacity=1, state_rows=state_rows)
94104
assert metadata.num_tokens.item() == 0
105+
assert not metadata.hidden.count_nonzero()
95106
assert metadata.indexed.positions.tolist() == [-1, -1, -1]
96107
assert metadata.indexed.token_slots.tolist() == [-1, -1, -1]
97108

98109

110+
@pytest.mark.parametrize("hidden", [torch.empty((2, 4096)), torch.empty((1, 4096), dtype=torch.bfloat16),
111+
torch.empty((4096, 2), dtype=torch.bfloat16).t(), torch.empty((2, 4096), device="meta", dtype=torch.bfloat16)])
112+
def test_draft_hidden_storage_rejects_incompatible_views(hidden):
113+
with pytest.raises(ValueError, match="draft hidden storage"):
114+
DraftMetadata(2, 1, 2, "cpu", hidden=hidden)
115+
116+
99117
@pytest.fixture(autouse=True)
100118
def records_only_cpu_commit(monkeypatch):
101119
# These state-machine tests run without CUDA; reduction and graph ordering

‎tests/recipes/test_glm53_stack.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,22 @@ def test_scratch_views_reject_unavailable_token_sample_and_request_capacity():
5959
base.rows(tokens, mode, sample_rows=rows)
6060

6161

62+
def test_cp_prefill_query_value_alias_survives_layer_overlay():
63+
base = DecoderScratch.allocate(65, 2, "prefill", "cpu", shard=TensorParallelShard(0, 2),
64+
sample_rows=1, context_ranks=2, overlay_layers=True)
65+
for rows in (65, 16, 65):
66+
view = base.rows(rows, "prefill", sample_rows=1).mla
67+
view.query.fill_(float("nan"))
68+
view.expanded.fill_(rows)
69+
assert torch.all(view.query == rows)
70+
for mode, ranks in (("prefill", 1), ("decode", 2)):
71+
separate = DecoderScratch.allocate(2, 2, mode, "cpu", shard=TensorParallelShard(0, 2),
72+
sample_rows=1, context_ranks=ranks).mla
73+
separate.query.fill_(7)
74+
separate.expanded.zero_()
75+
assert torch.all(separate.query == 7)
76+
77+
6278
@pytest.mark.parametrize("world,rank", [(2, 0), (4, 3)])
6379
def test_canonical_layers_share_intermediates_and_own_distinct_state(world, rank):
6480
snapshot = Path(os.environ.get("FLARE_GLM53_SNAPSHOT",

0 commit comments

Comments
 (0)