Repository navigation
Keep MTP on for GLM prompts past 2048 tokens - #398
Draft
vasilevklart wants to merge 50 commits into
Draft
vasilevklart wants to merge 50 commits into
vasilevklart wants to merge 50 commits into
Conversation
This was referenced Oct 11, 2026
vasilevklart
force-pushed
the
glm/09-mtp-longctx
branch
from
October 13, 2026 15:16
1a5fc68 to
a2db8c4
Compare
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.
vasilevklart
force-pushed
the
glm/09-mtp-longctx
branch
from
October 15, 2026 07:15
a2db8c4 to
4bcb6f5
Compare
2 tasks done
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.
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
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.
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.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, withKVManager.spare_rowas 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.max_tokensis 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 toindex_topk, so anymax_tokensafter a prompt past ~2k tokens was a 400 (every Claude Code request).M*:
e08def948(10-30); with MTP 1 server × 2 reps, without MTP 2 servers × 2 reps.Tests.
test_glm5next_mtp_cpuchecks that MTP equals greedy without MTP pastindex_topk, and each step's host rows at a real geometry. It serves joins, departures at their ownmax_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 istest_glm5next_mtp_graphon GPU.test_glm52_mtp_engine_cycle::test_mtp_stream_matches_baseline_past_topk_on_paged_dsadoes the same for GLM-5.2.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/rustand ruff pass.d877bfde4, job 507):test_glm5next_mtp_graphandtest_glm5next_mtp_cpupass (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.