Skip to content

Keep MTP on for GLM prompts past 2048 tokens - #398

Draft
vasilevklart wants to merge 50 commits into
glm/08-glm-chatfrom
glm/09-mtp-longctx
Draft

vasilevklart wants to merge 50 commits into
glm/08-glm-chatfrom
glm/09-mtp-longctx

Conversation

@vasilevklart

@vasilevklart vasilevklart commented Oct 10, 2026 •

Copy link
Copy Markdown
Collaborator

Problem. On #391 and #392, MTP turns off past index_topk, so any prompt over 2048 tokens decodes without drafts. An opencode session's first request is about 12k tokens.

Fix.

  • GLM-5.2 (configs/glm52_tp8_longctx_mtp.yaml): a verify step trims its rejected rows (latents and index keys) on the device before the next step plans.
  • GLM-5.3 (configs/glm5_next_tp8_longctx_mtp.yaml): no trims are needed. The next step's block overwrites the rejected rows, and every rank's page table moves the same way whatever was accepted. The DSA resource moves a verify block to the rows the device says it holds (DSAManager.move_rows, with KVManager.spare_row as the captured pad), so the model no longer re-addresses the KV plan or holds a view of the sink page. An eager prefill step seeds MTP in one pass over its last rows, not one per request.
  • A request's max_tokens is checked against what MTP decode can actually emit: the window, and the cache rows a verify step books a whole block of. Before, it was held to index_topk, so any max_tokens after a prompt past ~2k tokens was a 400 (every Claude Code request).
  • Greedy only, as with dense MTP. The trims, like the rest of GLM's MTP plumbing, stay model-local until Draft-verify as an engine resource for every MTP model #1195.
GLM-5.3, 8×H200, TP8, passkey, 256 tokens: decode ms/token 4k 32k 128k
M* with MTP 2.74 2.67 3.38
M* without MTP 4.88 4.92 5.15
vLLM 0.31 5.57 5.54 5.47

M*: e08def948 (10-30); with MTP 1 server × 2 reps, without MTP 2 servers × 2 reps.

Tests.

  • test_glm5next_mtp_cpu checks that MTP equals greedy without MTP past index_topk, and each step's host rows at a real geometry. It serves joins, departures at their own max_tokens, a reused slot and a stop token among accepted drafts as greedy, for MTP's own, always-right and half-wrong drafts at k = 1 and 3.
  • test_dsa_resource::test_move_rows_writes_and_attends_a_block_at_its_true_rows (its eager and captured addressing, on CPU); the captured attention is test_glm5next_mtp_graph on GPU.
  • test_glm52_mtp_engine_cycle::test_mtp_stream_matches_baseline_past_topk_on_paged_dsa does the same for GLM-5.2.
  • CPU Core on 3a614c4b8 (11,705 passed) fails no test that Serve GLM chat with roles, tool calls, reasoning and usage #397 (660e86f20) passes, by name; the TP lockstep sim, test/rust and ruff pass.
  • GPU, on an H200 at the stack's top with the upper fixes (d877bfde4, job 507): test_glm5next_mtp_graph and test_glm5next_mtp_cpu pass (54, 1 xfail). CPU Core is on torch 2.10, pyproject's floor; the GPU runs on 2.11.

Known limit. A GLM-5.3 verify step books its whole block and the host never takes the rejected rows back, so near the window admission refuses more than plain decode would (test_long_context_mtp_books_at_most_a_block_past_the_rows_it_holds, a strict xfail for the follow-up).

AI assistance (Claude) was used for this PR.

A verify step's next block overwrites its rejected rows on device: the host
plans whole blocks, as with holes, and the MTP layer selects with its own
indexer. Long prefill runs the MTP pass.
GLM-5.3's MTP rows read the forward's DSA context off the resource, and
the rule for decode rows trimmed on device (plan each at its longest)
moves with the step's staging into KPoolDSA.step_inputs.
GLM-5.3's MTP under long context lands on the merged DSA: the trimmed-row
lengths and the plans' masked padding (valid, pad_slot) move from kpool.py into
dsa/paged.py and manager.py, and its MTP index rows onto the aux planes.
# Conflicts:
#	mstar/model/glm5_next/submodules.py
test_long_context_mtp_output_equals_greedy_without_mtp is exact but
reaches ~130 rows: index_topk 64, pools of 4, windows of 48, two short
prompts. Nothing ran with the host's rows more than a page ahead of the
device's, or with hundreds of pools. The new case runs prompts of 700 and
450 tokens on 16-token pages with a 2048-token window (~175 pools against
the 16 a query selects) for 120 tokens at k = 3, with the MTP layer's own
drafts and with an oracle wrong every other draft, so the host counts
several pages of rejected rows the device rewrote. Every step checks the
tokens against greedy, the host's rows, the request's context and
MTP_LEN, as the short case does.
@vasilevklart vasilevklart changed the title GLM-5.2/GLM-5.3 PR 9: MTP with long context Keep MTP on for GLM prompts past 2048 tokens Oct 30, 2026
As glm52_tp8_longctx.yaml: mla_fused_prep on, which GLM-5.2's selecting
path takes since #392's _attend_written.
After the kernel cleanup on #388, glm52_tp8_longctx_mtp.yaml drops the switches it named;
kernels: auto serves what they turned on (MTP's prefill runs every row through the last layer,
as before).
MTP over DSA keeps only the decode graphs, so the three prefill bucket keys were ignored.
The models doc names both long-context MTP yamls.
… rows, not to index_topk

process_prompt held a max_tokens under MTP to index_topk - 2 (k + 1) + 1 - prompt, the
2048-token window without dsa_long_context. With it, a request past ~2,040 prompt tokens
that named a max_tokens was a 400 ("emits at most -17959 tokens"), which is every Claude
Code request (max_tokens 32000). The room is now what _mtp_check_stop leaves with no draft
accepted: the window less two blocks, and the cache rows at a whole block a step.
Brings main in through #397's merge. Conflicts:
- glm5_next/config.py: mla_cache_kpe keeps #398's docstring on #390's code,
  which drops _flashinfer_before with its FlashInfer 0.7.2 pin.
- glm5_next/submodules.py: #398's one window loop (MTP rows included), its
  uploads through #390's _upload.

Also, to fit main:
- PinnedStager has no to_device: the MTP seed's stored rows go through _upload.
- _mtp_seed writes MTP_LEN as int64: #389's scatter_ no longer casts
  (d59ef8f), and the DSA positions it starts from are int32.
Under mtp_trims the model re-addressed the KV plan's write slots itself, indexed the
DSA context's flat page list, and held a view of SINK_PAGE's last row in the KV
tensor (kv.kv_cache.tensor[:, SINK_PAGE, size - 1, ...]) to fill each captured step.
That is resource storage and plan state in model code (AGENTS.md invariant 1).

- KVManager.spare_row: a row no request holds (its flat slot and a view of it in
  every layer), MLA only.
- DSAManager.move_rows: moves the step's rows to device positions (write slots,
  positions; captured: per-row valid lengths and the spare row as the pad slot,
  masked with the caller's rope value; eager: host positions for an exact replan).
- The model computes MTP_LEN's positions and calls move_rows; no SINK_PAGE import,
  no tensor view, nothing bound at bind_node_resources.

Contract test: test_dsa_resource.py::test_move_rows_writes_and_attends_a_block_at_
its_true_rows (eager and captured branches on CPU). The long-context MTP parity
tests pass unchanged.
The windowed prefill re-ran each request's last MTP row in its own pass (one MTP
layer forward, with its TP all-reduces, per request in the step). One pass over
the step's last rows (DSAContext.pick, KVManager.token_rows) does the same work.

Unmeasured: the GPU stress plan times TTFT for many-prompt eager prefill steps.
The long-context MTP parity tests (2 prompts per prefill step) pass unchanged.
- glm5_next_tp8_longctx_mtp.yaml: the index planes are 3 x 128 wide, not 576, so
  the cache is ~23.0 KB per token per rank (~26.4 GB/rank), not 27.6 KB (~32 GB);
  and a prompt past 4096 tokens shares its eager step, as in glm5_next_tp8_longctx.yaml.
- glm52_tp8_longctx_mtp.yaml: only the MTP seed pass selects with its indexer; the
  draft passes attend densely (Glm52LLMSubmodule._mtp_decode).

Comments only.
… as greedy

The trims tests ran every request from prefill to the same max_tokens together.
Under mtp_trims each request's rows live in its KDA slot (MTP_LEN), so these serve
them as a worker does: prefills between decode steps, requests leaving at their own
max_tokens (cut mid-block), one taking a freed slot, a prompt that crosses
index_topk mid-decode and one that never reaches it, on 8-token pages. Each one's
tokens equal greedy's for MTP, always-right and half-wrong drafts at k = 1 and 3,
and MTP_LEN equals the rows held after every step. A stop token inside the
accepted drafts ends the reply where greedy's does.
Under mtp_trims the host books each verify step's whole block and never takes
the rejected rows back: with every draft rejected, a request's KV stream is at
108 rows after two steps while it holds 102. Pages run up to k + 1 times the
tokens kept, and process_prompt refuses a max_tokens past (kv_rows - prompt) /
(k + 1) that the window itself would hold.

A strict xfail until the follow-up: release step N - 1's rejected rows when
planning N + 1, from its verdict, which the host has by then and every rank
reads the same.
Brings in the base fixes, #391's/#392's DSA double-buffers, full-pool
xfail and sparse-plan bucket cap, and #397's refused-upload cleanup.
Conflicts:
- dsa/manager.py: #392's force_double_buffer property above #398's
  step_inputs signature (mtp_trims); both kept, they don't overlap.
- test_dsa_resource.py: both sides' new tests, #392's double-buffer test
  first, then #398's move_rows section.
It asserts that the prompt's MTP pass writes an index key for every
prompt row on the MTP layer's own plane; nothing in it checks what the
drafts select (parity with greedy holds whatever they select). The name
said it checked the selection.
vasilevklart added a commit that referenced this pull request Nov 5, 2026
Brings main in through #398's merge. Conflicts:
- serving_chat.py: imports. Main's LORA_ADAPTER_KWARG stays (its adapter
  pick in create_chat_completion merged clean); Utf8Stream goes with the
  code #813 moved to reply.py.

Also, to fit main:
- test_api_key's GLM server takes APIServer.lora_kwargs, which main's chat
  route calls (#820).
vasilevklart added a commit that referenced this pull request Nov 5, 2026
Brings in the base fixes and the #391-#398 fixes under #813. Conflicts:
- test_glm_chat.py: both sides' new tests at the same spot, #397's
  refused-upload test first, then #813's forced tool_choice tests.

This branch has not been deployed

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants