Repository navigation
Speed up GLM-5.2 with fp8 kernels and one-graph MTP - #388
Open
vasilevklart wants to merge 159 commits into
Open
vasilevklart wants to merge 159 commits into
vasilevklart wants to merge 159 commits into
Conversation
This was referenced Oct 9, 2026
vasilevklart
changed the title
This was referenced Oct 9, 2026
vasilevklart
marked this pull request as ready for review
October 9, 2026 21:55
vasilevklart
force-pushed
the
glm/03-glm52-speed
branch
from
October 10, 2026 22:20
c9d3d33 to
e9890df
Compare
vasilevklart
force-pushed
the
glm/03-glm52-speed
branch
from
October 13, 2026 15:16
e9890df to
7066ac6
Compare
GLM-5.2 at k=0: MLA with the DSA indexer, fp8 MoE, weight loader, configs. Replaces #272. Open: the configs' new max_concurrent_requests caps.
A prompt of limit - 1 tokens failed its whole decode batch (the speculative step overran the window), and a long-context prompt past index_topk failed its prefill batch. process_prompt now refuses both; prepare_inputs agrees.
Byte-level BPE models emit each token's raw bytes, which can end inside a character; decoding per chunk turned those into U+FFFD.
Per-token decode turned a character split across byte-level BPE tokens into U+FFFD; the 10-07 GSM8K answers have them at ÷ and √.
The conductor now builds it only for a model that sets prebuild_fused_moe.
As fused_experts_fp8: dynamo traced the launch path inline, which breaks the graph and fails on a symbolic token count under torch 2.9.1.
An fp8 checkpoint whose quant config went undetected loaded clean with unscaled e4m3 weights; a failed Hub fetch came back as a local path.
The runtime's ready order is handle order in Rust, and handles are recycled: a large prefill could be passed over indefinitely. Backlogged rows count too, since the batch builder tops a backlog up with fresh rows.
A non-numeric max_output_tokens from extra_body reached the conductor's min() and took its message batch down; image or audio parts were decoded and dropped. Both are 400s now.
Its logits kernel holds every token in one tile; past 128 the first launch raised OutOfResources. The models never send more; now it says so.
top_k and min_p reach the sampler; a fused-kernel import failure is logged; MSTAR_GLM52_GRAPH_COMPILE=0 also leaves uncaptured steps uncompiled; the single-rank test config is dummy mode, not a 750 GB pull.
Chosen per plan, a captured bucket could re-plan into the other kernel while the graph replayed the captured one on stale buffers.
from_hf_config read only rope_theta, so a YaRN or dynamic variant ran as plain RoPE, its frequencies and softmax scale wrong.
The CUDA op launched a 0-block grid unchecked; the latched error surfaced in the next unrelated kernel.
Each new test file with a CPU case, listed in test/cpu-core.txt.
A rank whose setup failed fell back to NCCL alone, and its peers' first all-reduce waited on it forever. tp_allreduce_max_kb: 0 now turns it off instead of reading as unset (512).
Built from tokenizer.json alone (transformers 4.x), the tokenizer had no special tokens, so the detokenizer printed <|user|> at the end of replies.
components/mla.py ran its fused prep (mla_prep: q_a's norm, then kv_a's norm, both RoPEs and the latent's cache write, in two kernels) only for a forward that selects no rows. A selecting subclass that attends rows the prep already wrote (_attends_written, _attend_written) now takes it too: q_norm first, the selection from its q_c, then kv_write_q_rope and attention over the selected slots. The kernels reproduce the unfused path bit for bit. Nothing here sets _attends_written, so every model keeps its path on this branch; GLM-5.2's selecting attention takes it with #392's paged DSA.
decode.py's H200 table was swept at GLM-5.3's hidden 4096, and at
GLM-5.2's experts (hidden 6144, 256 experts of 256 at TP8) it is the slow
one at 1-6 and 9-10 tokens. A table swept at that shape on an H200, with
every call's weights read from HBM (the per-pair winner per token count up
to 8, the grouped kernels from 9), now serves that shape on every GPU but
the H100; any other shape keeps the H200 table, and the H100 keeps its own.
One layer's experts (routed and shared), CUDA-graph replay, us a call,
H200 table -> GLM-5.2's:
tokens 1 2 4 5 6 8 9 10
29.85 33.38 53.52 70.87 81.02 90.45 113.13 116.85
-> 24.15 31.12 51.94 63.94 75.37 89.78 102.08 107.80
GLM-5.2-FP8 on 8x H200 (TP8, glm52_tp8.yaml), random 1024 in / 512 out,
ignore_eos, greedy, a fresh server per arm: c=1 131.3 -> 138.7 tok/s
(TPOT 7.48 -> 7.06 ms), c=8 555.1 -> 556.6. GSM8K-200 183 -> 184. Greedy
parity (8 prompts x 128 tokens): 3 of 8 identical, the others first
differing at tokens 7 to 71, since the new tiles sum in another order; a
second server whose decode MoE path was unchanged matched 8 of 8.
Tests: test_fused_moe_decode.py (the table pick by GPU name and shape).
…3 tokens on the H200 On an H200, FlashInfer's SM90 block-scale GEMM beats the w8a16 kernel from 3 tokens at GLM-5.2's o_proj (N 6144, K 2048) and dense-layer down (N 6144, K 1536), where the shape table kept w8a16 up to 32 tokens. Cold weights on one H200, us a call at 4 / 8 / 16 / 32 tokens: o_proj w8a16 10.32 / 11.05 / 9.84 / 10.59 FlashInfer 8.01 / 7.99 / 8.09 / 9.28 dense down w8a16 9.27 / 9.28 / 8.74 / 9.55 FlashInfer 7.51 / 7.56 / 7.63 / 7.97 The new rows are keyed by GPU model, since the H100 and the H200 are both SM90. reserve() resolves a device's limits once, at load and outside capture. They apply only where FlashInfer's GEMM runs, so a GPU without it keeps w8a16 at these sizes instead of falling through to cuBLAS's slower GEMM. Every other GPU and shape keeps its path. A 1x H100 shows the same crossover (o_proj 10.53 -> 9.04 us at 3 tokens, 11.65 -> 9.09 at 8), but GLM-5.2-FP8 does not fit 8x H100, so the H100 gets no row until it is measured end to end. FlashInfer quantizes the activations to e4m3 per 128 columns, as dense_fp8 already does from 33 tokens, so outputs move at 3-32 tokens. GLM-5.2-FP8 on 8x H200 (TP8, configs/glm52_tp8.yaml; random 1024 in / 512 out, greedy, closed loop; a fresh server per arm, in the order base, this, this, base), tok/s: c base this change TPOT p50 (ms), mean of the rounds 1 138.7 / 138.7 138.8 / 138.6 0.0% 7.06 -> 7.06 8 556.6 / 556.8 564.2 / 563.1 +1.3% 13.76 -> 13.50 32 1177.6 / 1270.8 1260.7 / 1280.6 +0.8% 23.42 -> 23.36 (second round) 64 1924.7 / 1963.6 1919.1 / 1968.5 0.0% 29.82 -> 29.65 At c=32 both first-round servers stalled for up to 2.4 s, the base's more, so only the second round is compared there. c=1 and c=64 are controls: at 1 token nothing changes, and from 33 tokens both arms already run FlashInfer's GEMM. With MTP (one server each), c=1 goes 197.3 -> 199.9 tok/s (+1.3%) and c=8 785.2 -> 792.1 (+0.9%). GSM8K-200 is 185 against 183, with 0 errors on both. Tests: test_glm52_kernel_platform.py checks that the H200 takes FlashInfer from 3 tokens on these shapes, and that an H100, an H200 NVL and an H200 without FlashInfer keep w8a16 up to 32 tokens. test_glm52_dense_fp8.py checks reserve()'s limits, and that FlashInfer stays within the rounding bound at 3 and 8 tokens.
components/mla.py carried three options for Kimi K3, which no GLM model sets: an output gate (mla_use_output_gate: g_proj and _gate), the latent norms' own eps (mla_latent_norm_eps), and unrotated decoupled dims (mla_use_nope with qk_rope_head_dim > 0). GLM-5.2 rotates its decoupled dims and GLM-5.3 has none; K3 is not in this launch, so the options and their tests move to the K3 stack (refs/prep/k3/00-mla-k3-options). mla_use_nope stays for GLM-5.3 and now means no decoupled dims: with qk_rope_head_dim above 0 it is refused. Every GLM path is unchanged: with the gate off _gate returned its input, the eps fell back to rms_norm_eps, and GLM-5.3's config already refuses mla_use_nope with rope dims. Tests: test_mla_variants.py keeps GLM-5.3's layout and pins the new refusal; K3's reference and gate tests leave with the options.
… the checkpoint GLM-5.2 had nine kernel switches (moe_quant_kernel, moe_decode_kernel, moe_router_kernel, moe_prefill_kernel, moe_fused_allreduce, dense_fp8, mla_fused_prep, fused_add_rmsnorm, prefill_last_layer_rows), off by default and on in every serving yaml. Now: - The MoE prefill kernels, which keep the W8A8 runner's tiles and K order (the same bits, test_glm52_moe_decode.py), run on every fused block, with no switch. - The fp8 linears follow the checkpoint: the non-expert linears a block fp8 checkpoint stores in fp8 stay fp8 (dense_fp8 is now a property of the quantization config), and a bf16 config, as the reduced test models are, takes the bf16 path. - What changes numerics is one option, model_kwargs.kernels. "auto" (the default) is what the yamls turned on: the fused fp8 MoE with fused_moe's decode chain and router kernel (their per-GPU tables live in fused_moe/decode.py and dense_fp8.py), the MLA latent prep as two kernels, each all-reduce fused with the next add + RMSNorm and the MoE's partials reduced once, and the last layer's FFN on the sampled rows at prefill. "reference" is the numerics anchor, served eager: the per-expert dequant loop, the fp8 linears dequantized per call, the latent prep op by op, unfused all-reduces and norms, every prefill row through the last layer. "triton" (the fused MoE or an error) is gone: auto serves the reference loop, with a warning, where the fused kernel can't run. The fused latent prep sits under auto, not beside the prefill kernels: its kernels have the ops' bits, and plain decode keeps them (a reduced GLM-5.2 on an H200, round 9 without it against this with it: 3/3 plain cases identical), but 3 of 7 MTP cases differ, in KV pages (two captured cases) and in one eager case's LM-head bits, with every token the same. A config that still names a retired switch is refused at startup, and a fp8 linear the checkpoint doesn't store is named by the loader's missing-weight check. The serving yamls drop the switches and serve what they served: GLM-5.2 bitwise, round 9 with the yamls' switches against this with auto, 10/10 identical; round 9's defaults against this with reference, 10/10. Every other model is unchanged: only GLM-5.2's config names kernels. Tests: the GLM-5.2 kernel, MoE, fused-norm, MLA-prep and dense fp8 files ask for kernels instead of the switches; the integration checkpoint stores the attention, shared-expert and dense-MLP linears in fp8 as the real one does.
Comments and docstrings keep the why and lose the history. Tests that re-checked a switch, a default or a log line, or copied production logic (the unused host verify helpers, the loader's key inventory) are cut or merged into parametrized cases.
…d verify helpers The two custom ops live in mstar/utils and serve GLM-5.3 too. mtp_greedy_verify and its host twin had no caller outside tests: verify runs on the device in the step. The models doc names glm52_tp8_mtp.yaml.
…piled calls GLM-5.2's submodule raised torch._dynamo.config.recompile_limit for the whole thread (the whole process on torch < 2.12), so every other node on a worker compiled under it. A submodule now asks through NodeSubmodule.dynamo_recompile_limit(), and the engine and the graph runner patch the limit around that submodule's compiled callables only. One that doesn't override the hook gets its callables back unwrapped.
dghosef
requested changes
Nov 3, 2026
dghosef
left a comment
dghosef
left a comment
Collaborator
There was a problem hiding this comment.
Ty - a few things in my comments, mainly the one-rank verdict timeout, the fused MLA prep's kv_label and rope_interleave gaps, and measuring the engine changes on other models.
| # prepare_inputs runs, so the last verdict's trim lands here. The verify and the | ||
| # drafts are argmaxes in the forward: no sampler. | ||
| for rid in request_ids: | ||
| self._apply_verdict(rid) |
Collaborator
There was a problem hiding this comment.
Just make sure that once we get the actual spec decoding resource in (#1090), you move this onto it. Right now the MTP state is 8 per-request dicts and a seed slot pool in the model, the host spin-waits for each step's verdict outside check_stop (so MTP steps lose the async pre-plan overlap), and rejected rows are taken back through the new correct_len methods on KVManager/RopeManager. When that lands, can you delete all of this, correct_len included?
| return dense_fp8.linear(x, weight, scale, block_size, glu) | ||
|
|
||
|
|
||
| class Fp8Linear(nn.Module): |
Collaborator
There was a problem hiding this comment.
Non-blocking, just fyi - Fp8Linear re-does the column/row sharding and the reduce that ColumnParallelLinear/RowParallelLinear already do, and each call site picks one w/ a ternary. It also isn't LoraCapable, so LoRA can't target q_b_proj/o_proj on fp8 checkpoints. Longer term it'd be nice if the weight format were a quant argument on the existing parallel linears.
| hidden_states = residual + hidden_states | ||
| return hidden_states + self.mlp(self.post_attention_layernorm(hidden_states)) | ||
|
|
||
| def _forward_fused(self, hidden_states, residual, position_ids, rows): |
Collaborator
There was a problem hiding this comment.
This loop is a second copy of run_layers / forward_mlp_residual in mstar/model/components/decoder_layer.py, so a fix to one won't reach the other, and only GLM gets the sampled-rows trick. Can you open an issue for moving this onto run_layers (w/ position ids and a last-layer rows option)?
…ops the worker A rank that timed out waiting for a verdict, or found its mailbox reused, raised RuntimeError on its own: its batch failed while the other ranks went on into the next all-reduce and hung. Every read follows the step's launch (a pre-plan waits on the step's commit), so only a hung device or peer delays a verdict, as with any collective. The overwrite check is a broken invariant rather than a request's failure, so it raises SystemExit, as micro_scheduler does for diverged TP ranks.
_write_fused_and_attend called write_slots() and the attention's run() without the label, so a RoPE layer with kv_label set wrote and read the default stream on the fused path. Both paths take it from one property now.
mla_prep's _rope turns the pairs (2i, 2i + 1), while the op-by-op path follows PositionConfig.interleave, so a checkpoint with "rope_interleave": false under kernels: auto had q_pe and k_pe rotated on the wrong pairs with no error. RopeManager exposes interleave, and bind_resources keeps the ops for NeoX pairing, as it does without a cos/sin table.
…a loop The MoE prefill kernels run only on a fused block, and only kernels: auto builds one (Glm52SparseMoeBlock.process_weights_after_loading).
…2.12 dynamo's config.patch is process-wide there, and pyproject still allows torch 2.10; the limit stays per submodule only because a worker runs one compiled call at a time.
Glm52SparseMoeBlock set NoAuxTcRouter._topk_op from outside. Whether the kernel runs is known only on the loaded device, so the router's own load hook takes it as an argument and imports the op. Glm52LLMSubmodule reads its config, and an attribute __init__ sets, instead of getattr fallbacks, and the GLM-5.2 MoE block reads config.kernels.
get_request_resource_configs built the sampler config from a fixed key list without penalize_prompt, so a request that set it false with a repetition_penalty still had its prompt penalized. It now passes with top_k and min_p as a sampler knob with no config default (as #390 does for GLM-5.3). test_glm52_request_kwargs: every knob the sampler declares (REAL_KNOB_RULES, INT_KNOB_RULES, BOOL_KNOBS) reaches the sampler config or is refused at intake.
vasilevklart
added a commit
that referenced
this pull request
Nov 3, 2026
…t stops the worker The same fix as GLM-5.2's (#388). A rank that timed out waiting for a verdict, or found its mailbox reused, raised RuntimeError on its own: its batch failed while the other ranks went on into the next all-reduce and hung. Every read follows the step's launch (a pre-plan waits on the step's commit), so only a hung device or peer delays a verdict, as with any collective. The overwrite check is a broken invariant rather than a request's failure, so it raises SystemExit, as micro_scheduler does for diverged TP ranks.
Conflicts: - kv/manager.py: main's plan states (#932, #1127): every eager packed step's addressing is built on the host, and decode takes the resident rows. #388's small-step host path and its thresholds go, with KVPlanOutput.max_span and test_kv_host_plan.py. correct_len now uses tokens_per_page, cache.copy_pages and StreamPages' del/extend, and refuses a cache with KV groups. - glm52/submodules.py: main's TokenStop (#794) is the plain decode's stop. Under MTP, check_stop counts a verify step's tokens, and a prefill stop or one-token budget returns END_PARTITION, which replaces the ended_at_prefill latch. The prefill groups read PREFILL_TOKEN_BUCKETS. - fused_moe/decode.py: main's kernel registry (#986). GLM-5.2's H200 tiles move from _H200_GLM52 to kernels/tables/moe_decode/NVIDIA_H200.json, so they now apply on H200 only. - mla.py, adding_models.rst: both sides' additions. - test_glm52_dsa_engine.py, test_fused_moe_decode.py: #388's tests on the merged code. Also, to fit main: - GLM-5.2's per-step index upload goes through to_device_nonblocking, since PinnedStager no longer has to_device. - prefill_token_buckets is a named Glm52Model argument, which the shipped-ladder test checks. - #388's tests take main's capture_into_graph stream, the arena's owners/is_sealed, and swap the KV manager's arena (now a read-only Rust object) rather than patch it. - test_glm52_capture_gpu's harness ends decode at Glm52Model.get_loop_caps, as the runtime does: since #857 check_stop stops on EOS only, so main's own copy of this test streams 15 tokens for 12 and fails.
vasilevklart
added a commit
that referenced
this pull request
Nov 3, 2026
Brings main in through #388's merge. Conflicts: - recurrent/pool.py: main's one admit critical section, version bumps and per-dtype _foreach_zero_ cache (_slot_views), with #389's pre-planned reservation guard, the _release helper (now bumping the version) and per-slot blocks; #389's slot generator is renamed _slot_rows. - runner.py: main's own refused-admit stage drop, which is #389's fix. - convenience.py: main deletes RaggedAttentionCallable (#858); KDACallable stays. - rid_table.py, adding_models.rst, test_admit_rollback.py: both sides. Also, to fit main: - KDA's eager layout upload copies through its stager into a fresh tensor, since PinnedStager no longer has to_device. - #389's tests ingest a request once (main's worker does, #725). - test_worker_node_metadata runs on the Rust graph runtime, main's only one, and stubs what main's Worker.__init__ reads. - main's test_chunked_prefill engine fake sets _tp_follower_nodes, which the follower-refusal exit reads.
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. GLM-5.2 on main (#387) dequantizes its fp8 linears on every call and routes the MoE through the per-expert reference loop, which can't be captured, so decode runs eager and the model has no MTP.
Fix.
kernels: auto, the default, keeps the fp8 linears fp8: w8a16 at decode sizes, FlashInfer's block-scale GEMM above. The same mode adds:kernels: referenceis the op-by-op numerics anchor.configs/glm52_tp8_mtp.yamladds greedy MTP k=3 as one captured graph per batch size; its plumbing stays model-local until a follow-up moves it into a shared draft-verify resource.components/mla.pybecomes the MLA layer GLM-5.3 reuses (#390). A submodule's dynamo recompile limit (dynamo_recompile_limit()) now applies around its own compiled calls only.1024 in / 512 out, greedy, a fresh server per arm; vLLM 0.31 and SGLang 0.5.21 at TP8, 2048 context, bf16 KV, k=3.
Tests.
test_glm52_mtp_engine_cycle: the MTP stream equals plain decode bitwise.test_glm52_fused_add_rmsnorm: fused wiring emits what unfused does.test_glm52_moe_decode(GPU): the H200 tile table against the runner.517d41d65(8,846 passed) fails no test that main (9f84f18cd) passes, by name; ruff is clean.91f303959) fail only the 22 tests that need an HF cache. CPU Core is on torch 2.10, pyproject's floor; the GPU runs on 2.12.1.Not run yet. Qwen3.5-0.8B at c=32/64 against main, for this PR's engine changes (host KV plan, non-blocking copies, padding release).
AI assistance (Claude) was used for this PR.