Repository navigation
Share the first GLM draft pass's sparse selection with later passes - #218
original-el8 wants to merge 1 commit into
Conversation
The GLM-5.3-Flash checkpoints declare `index_share_for_mtp_iteration: true`: every draft pass of a speculative round attends through the sparse selection of the round's first pass. Transformers has no GLM MTP layer and ignores the field. vLLM's GLM MTP proposer runs the draft indexer at the first draft step and reuses its top-k selection for later steps when prefill context parallelism is off. Flare's draft passes ran the MTP layer's indexer, and with context parallelism its candidate exchange, on every pass. When the checkpoint configuration declares the field true, rounds with two or more drafts run the MTP indexer only on their first draft pass and copy each request's physical selection and selected length into round-local buffers. Before each later pass, every proposing request's selection is placed in its packed draft row, and idle rows select nothing. Later passes skip the index query and gate projections, query quantization, C4 pool update, selection, expansion, and context-parallel candidate all-gather. They still write their latent records, and they attend to records at or before the first pass's position rather than to the records of the round's earlier drafts. Only the first pass advances the draft C4 tail ring, which the round restores before absorption. Target verification and absorption are unchanged, so the served distribution is unchanged; draft proposals and acceptance follow the declared selection semantics. Checkpoints that do not declare the field run the indexer on every pass. docs/mtp-design.md describes the shared selection. Compatibility: no package metadata or prepared-contract change. Exported packages keep running the indexer on every draft pass; re-export prepares the shared-selection draft graphs. Validation on main (972e28e) with this change, on one RTX PRO 6000 (SM120): - tests/recipes/test_glm53_mla_cuda.py: test_reused_selection_attends_only_to_the_first_pass_records passed in 4 cases (TP2 and TP4 head partitions, query-chain overlap on and off) and skipped 2 (TP1), with layer 3 MLA weights of GLM-5.3-Flash-NVFP4-MXFP8-CSF-QAD. With 256 cached records, a first pass at position 256 selects 257 records. A reusing pass at position 257 leaves the selection, C4 cache and C4 tail unchanged, writes its latent record, and produces zero output for the idle row. Its output matches a fresh selection at position 256 within 1e-3 relative norm (bitwise in all four cases) and differs by 11 to 17 percent relative norm from a selection at its own position. The file's six TP1 native-replay cases fail during b12x BMM tuning ("BMM plan requires a supported BF16 x MXFP8 geometry"), identically without this change. - tests/recipes/test_glm53_speculative.py with the other GLM recipe, speculative-capacity and prepared-contract CPU tests: 223 passed, 47 skipped. When a request stops proposing and the packed rows shift, each remaining request receives its own first-pass selection. 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. - 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. Limitations: under decode stream overlap the indexer runs beside the auxiliary-stream query projection and absorption, so skipping it on a later pass saves less than its duration. With context parallelism, each rank reuses its own first-pass local selection; that path is untested. Not run: context parallelism, SM103 hardware, and acceptance or decode-time measurements that isolate this change. Design-Signed-Off-By: flare-02; design-review=4d3a
|
Categorized as Review request-to-packed-row ownership, idle rows, C4 restoration, and alternating round widths together with #217. The reported single-GPU partition tests and combined multi-host serving runs leave context-parallel selection reuse untested. Qualify that affected path and measure acceptance/round time on a controlled comparison before attributing a speedup to this PR. The implementation is not already present on main at 972e28e. |
Share the first GLM draft pass's sparse selection with later passes Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Design-Signed-Off-By: Claude Fable 5.1 (Claude Code); design-review=4d3a Signed-off-by: Kyonjova <kyonjova@pm.me>
|
@original-el8 GitHub reports merge conflicts with
Please reconcile the branch with main and rerun the checks affected by the resolution. |
The GLM-5.3-Flash checkpoints declare
index_share_for_mtp_iteration: true:every draft pass of a speculative round attends through the sparse selection
of the round's first pass. Transformers has no GLM MTP layer and ignores the
field. vLLM's GLM MTP proposer runs the draft indexer at the first draft step
and reuses its top-k selection for later steps when prefill context
parallelism is off. Flare's draft passes ran the MTP layer's indexer, and
with context parallelism its candidate exchange, on every pass.
When the checkpoint configuration declares the field true, rounds with two
or more drafts run the MTP indexer only on their first draft pass and copy
each request's physical selection and selected length into round-local
buffers. Before each later pass, every proposing request's selection is
placed in its packed draft row, and idle rows select nothing. Later passes
skip the index query and gate projections, query quantization, C4 pool
update, selection, expansion, and context-parallel candidate all-gather.
They still write their latent records, and they attend to records at or
before the first pass's position rather than to the records of the round's
earlier drafts. Only the first pass advances the draft C4 tail ring, which
the round restores before absorption. Target verification and absorption are
unchanged, so the served distribution is unchanged; draft proposals and
acceptance follow the declared selection semantics. Checkpoints that do not
declare the field run the indexer on every pass. docs/mtp-design.md
describes the shared selection.
Compatibility: no package metadata or prepared-contract change. Exported
packages keep running the indexer on every draft pass; re-export prepares
the shared-selection draft graphs.
Validation on main (972e28e) with this change, on one RTX PRO 6000 (SM120):
test_reused_selection_attends_only_to_the_first_pass_records passed in
4 cases (TP2 and TP4 head partitions, query-chain overlap on and off) and
skipped 2 (TP1), with layer 3 MLA weights of
GLM-5.3-Flash-NVFP4-MXFP8-CSF-QAD. With 256 cached records, a first pass at
position 256 selects 257 records. A reusing pass at position 257 leaves the
selection, C4 cache and C4 tail unchanged, writes its latent record, and
produces zero output for the idle row. Its output matches a fresh
selection at position 256 within 1e-3 relative norm (bitwise in all four
cases) and differs by 11 to 17 percent relative norm from a selection at
its own position. The file's six TP1 native-replay cases fail during b12x
BMM tuning ("BMM plan requires a supported BF16 x MXFP8 geometry"),
identically without this change.
speculative-capacity and prepared-contract CPU tests: 223 passed,
47 skipped. When a request stops proposing and the packed rows shift, each
remaining request receives its own first-pass selection. 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.
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.
Limitations: under decode stream overlap the indexer runs beside the
auxiliary-stream query projection and absorption, so skipping it on a later
pass saves less than its duration. With context parallelism, each rank
reuses its own first-pass local selection; that path is untested. Not run:
context parallelism, SM103 hardware, and acceptance or decode-time
measurements that isolate this change.