Skip to content

Share the first GLM draft pass's sparse selection with later passes - #218

Open
original-el8 wants to merge 1 commit into
mainfrom
perf/glm53-draft-selection-sharing
Open

original-el8 wants to merge 1 commit into
mainfrom
perf/glm53-draft-selection-sharing

Conversation

@original-el8

Copy link
Copy Markdown
Collaborator

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.

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
@original-el8 original-el8 added area:kernels GPU kernels and numerical execution. area:preparation Ignition checkpoint loading, tuning, compilation, capture, and export. bug Something isn't working performance Performance, latency, throughput, or memory-efficiency investigation or improvement. priority:P2 Normal bug priority: workflow failure or diagnosis with narrower impact or an available workaround. triage:review Proposed change needs code review; this does not imply merge readiness. labels Nov 6, 2026
@original-el8

Copy link
Copy Markdown
Collaborator Author

Categorized as triage:review, bug, priority:P2, and performance, under kernels and preparation. This honors the checkpoint's index_share_for_mtp_iteration declaration and avoids repeating the draft indexer; it changes draft proposal semantics while retaining target verification.

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.

kyonjova added a commit to kyonjova/flare-local-inference that referenced this pull request Nov 8, 2026
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 original-el8 added the model:glm-5.3 Affected recipe or documented reproducer: glm-5.3. Impact may be broader. label Nov 8, 2026
@original-el8

Copy link
Copy Markdown
Collaborator Author

@original-el8 GitHub reports merge conflicts with main, confirmed locally with git merge-tree for PR head 2f99faf20087 against main 97cfeb0001f1. Conflicting paths:

  • ignition/recipes/glm53/mla.py
  • ignition/recipes/glm53/mtp.py
  • ignition/recipes/glm53/speculative.py
  • tests/recipes/test_glm53_mla_cuda.py
  • tests/recipes/test_glm53_speculative.py

Please reconcile the branch with main and rerun the checks affected by the resolution.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:kernels GPU kernels and numerical execution. area:preparation Ignition checkpoint loading, tuning, compilation, capture, and export. bug Something isn't working model:glm-5.3 Affected recipe or documented reproducer: glm-5.3. Impact may be broader. performance Performance, latency, throughput, or memory-efficiency investigation or improvement. priority:P2 Normal bug priority: workflow failure or diagnosis with narrower impact or an available workaround. triage:review Proposed change needs code review; this does not imply merge readiness.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant