Repository navigation
[None][perf] GLM-5.3-Flash FP8 performance fixes on Hopper - #21179
Draft
ruocheng-nv wants to merge 8 commits into
Draft
ruocheng-nv wants to merge 8 commits into
ruocheng-nv wants to merge 8 commits into
Conversation
Without tcgen05 the fused_hc autotuner only had the FP32 FMA kernels, which stop scaling with the token count: a 13K-token prefill spent 2-3 ms per mHC boundary on H200. Add a "split_dg" fused_hc backend for pre-SM100 GPUs and M > 32: the previous block's post_mapping, then pre_mapping's DeepGEMM TF32 GEMM + sqrsum (SM90 tensor cores, optional split-K) and big_fuse. The autotuner still picks the fused FMA kernels for decode-sized M. fused_hc only tunes token buckets up to 8192, so a larger prefill chunk runs the fallback tactic; on pre-SM100 GPUs that fallback is now split_dg for M > 32 when DeepGEMM is available. The lagged pre-mix mode (pre_mix_ext) no longer pins the pre_mapping and fused_hc boundaries to the FMA kernels: big_fuse already carries the lagged variant for every split count, so the DeepGEMM tactics serve it too. H200, n=4, hidden 4096, M=13312 (best FMA tactic -> best DeepGEMM tactic): fused_hc 1717 -> 535 us (V4 and lagged) pre_mapping lagged 2471 -> 304 us Differences vs the FMA path are at TF32 level (layer_input within one bf16 ulp). GLM-5.3-Flash FP8, H200 TP4, 13K-token prompt, concurrency 1: TTFT 857 -> 645 ms. Signed-off-by: Ruocheng Jia <ruochengj@nvidia.com>
FlashMLA's SM90 sparse prefill kernel walks the top-k axis two 64-wide blocks at a time and asserts topk % 128 == 0. glm_kpool padded its index rows to a multiple of 64 only (2051 -> 2112), so GLM-5.3-Flash failed during warmup on Hopper with "Assertion `params.topk % (2*B_TOPK) == 0`". Pad the expanded index rows to 128 on SM90 (2176) and keep 64 elsewhere, so the SM100 native decode path and its 2112-wide rows are unchanged. Signed-off-by: Ruocheng Jia <ruochengj@nvidia.com>
On SM90 the GLM-5.3-Flash sparse attention runs FlashMLA's sparse prefill kernel, which processes one query token per CTA and walks all of its top-k rows (2176 with the k-pool tail). A decode or MTP-verify call with a few tokens therefore occupied a handful of SMs for ~160 us per layer. For at most 16 query tokens on SM90, the call now splits each token's top-k rows into 128-row chunks, runs one kernel row per (token, chunk) and merges the chunk outputs by their log-sum-exp in a small Triton kernel (chunks without valid rows contribute nothing). H200, 16 local heads, 2176 top-k rows: 1/4/8 tokens 165/162/167 us -> 30/36/48 us. GLM-5.3-Flash FP8, H200 TP4, MTP-3, concurrency 1: 297 -> 316 output tok/s per user. GSM8K unchanged. Signed-off-by: Ruocheng Jia <ruochengj@nvidia.com>
Where the in-tree KDA prefill op is unavailable (e.g. Hopper), prefill runs FLA's chunk_kda, which derives its chunk indices from cu_seqlens. Given only the device tensor, that repeat_interleave reads the chunk counts back from the GPU, so every KDA layer stalled the host until the GPU drained. The Kimi K3 mamba metadata already has the context lengths on the host; it now also keeps a host cu_seqlens, and the mixer passes it to chunk_kda as cu_seqlens_cpu so the indices are built on the CPU. The metadata also keeps one device cu_seqlens view per forward: FLA caches its chunk metadata by tensor identity, and a fresh slice per layer made every layer rebuild and re-upload it. GLM-5.3-Flash FP8, H200 TP4, 13K-token prefill (34 KDA layers), concurrency 1: the host cu_seqlens took the prefill step from 486 to 474 ms (TTFT 507 -> 500 ms); the shared device view, measured on a later build, a further 462 -> 459 ms (TTFT 489 -> 483 ms). Signed-off-by: Ruocheng Jia <ruochengj@nvidia.com>
Without the fused verify kernel (e.g. on Hopper), KDA target verification looped over the draft steps: per step three short convolutions, one FLA recurrent call, and copies of the step's conv windows and state into the intermediate buffers. With MTP-3 that is about 40 launches per KDA layer. When the packed convolution weight is available and the intermediate state buffer is fp32 (or the pool's dtype), all steps now run as one packed causal_conv1d_update that writes the per-step windows into intermediate_conv_window, and one fused_recurrent_kda call that seeds row (b, 0) of intermediate_ssm with the pool state and writes step t's state to row (b, t). The convolution runs on a gathered copy of the rows, so the live pools stay read-only as before; the step index tensors are built once per shape and Q/K/V are read in place from the fused projection. The per-step loop remains the fallback; the verify parity test now covers both paths. GLM-5.3-Flash FP8, H200 TP4, MTP-3, 13K-token prompt: decode step 12.1 -> 9.6 ms; concurrency 1: 226 -> 274 output tok/s per user; concurrency 8: 59.8 -> 65.1 output tok/s. GSM8K unchanged. Signed-off-by: Ruocheng Jia <ruochengj@nvidia.com>
…kernel finalize_decode_weights returned early unless the optimized KDA decode kernel was in use, so without it (e.g. on Hopper) the q/k/v projections and the [f_a | g_a | b] projections ran as separate GEMMs in every prefill, decode and verify step: nine small GEMMs per KDA layer per decode step instead of four. The projection fusion now runs for every decode kernel path; only the optimized kernel's layout constants stay conditional. GLM-5.3-Flash FP8, H200 TP4, MTP-3, 13K-token prompt, concurrency 1: 274 -> 316 output tok/s per user, TTFT 500 -> 489 ms; concurrency 8: 65.0 -> 67.2 output tok/s. GSM8K unchanged. Signed-off-by: Ruocheng Jia <ruochengj@nvidia.com>
DeepGEMM ships an SM90 fp8_mqa_logits kernel, so the FP8 context pool scorer is not SM100-specific. Enable it on SM90 as well. On H200 at a 13K-token prompt, pool scoring drops from 1.20 ms to 0.35 ms per sparse layer (top-512 selections overlap 98% with the BF16 scorer). End to end with TP4 and MTP3 (ISL 13312 / OSL 33): C1 TTFT 483 -> 475 ms; C8 output throughput 68.2 -> 69.7 tok/s. GSM8K (400 questions): 98.5%. TRTLLM_GLM_INDEXER_BF16=1 still selects the BF16 scorer. Signed-off-by: Ruocheng Jia <ruochengj@nvidia.com>
Short prompts are host-bound in eager prefill. Make the model work with prefill_cuda_graph_backend: breakable (with disable_mm_encoder: true for text serving): - KDA layers join the shared mixer's breakable-graph registry, so their metadata-dependent core runs between graph segments. - Sparse attention runs its phase split as one eager op between graph segments on the live rows and zeroes the padded rows; the TP reduction stays in the graph. - Eager passes over a padded batch (warmup) run only the live rows. H200 TP4 MTP3, C1: ISL 1K TTFT 130 -> 86 ms; 4K and 13K unchanged. GSM8K (400 questions) 98.0%. Signed-off-by: Ruocheng Jia <ruochengj@nvidia.com>
ruocheng-nv
force-pushed
the
perf/glm53-flash-hopper-py
branch
from
October 29, 2026 06:39
9b430ab to
7bc696d
Compare
This was referenced Oct 29, 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.
Description
Performance fixes for the GLM-5.3-Flash FP8 checkpoint on Hopper (H200, SM90), measured at TP4 with MTP3 and BF16 KV
cache on a long-input workload (ISL 13312 / OSL 33). Split out of #20878:
-rdc, CUTLASS FP8 block-scale MoE row expansion and activation): [None][perf] Build FlashMLA without -rdc and speed up FP8 block-scale MoE expansion/activation #21178.This PR does not depend on either for correctness.
Attention / indexer
at a time, and GLM-5.3-Flash asserted during warmup on Hopper.
Triton kernel, so small decode batches fill the GPU.
fp8_mqa_logits, added for SM100 in [None][perf] Speed up GLM-5.3-Flash serving on Blackwell #20755) on SM90.TRTLLM_GLM_INDEXER_BF16=1keeps the BF16 scorer.mHC
buckets). The lagged pre-mix mode no longer pins the boundaries to the FMA kernels, so the pre-mapping autotuner also
offers DeepGEMM tactics for the lagged mode on SM100.
KDA (shared Kimi KDA mixer: applies to every model on its FLA / non-optimized paths)
cu_seqlensto FLA chunked prefill so it does not sync the device per layer.Breakable prefill CUDA graphs
prefill_cuda_graph_backend: breakable(setdisable_mm_encoder: truefor textserving). KDA layers join the shared mixer's breakable-graph registry. Sparse attention runs its phase split as one
eager op between graph segments on the live rows. Eager passes over a padded batch run only the live rows. Buckets
above ~2K tokens are not recommended.
Performance (H200 x4, TP4, MTP3, BF16 KV, ISL 13312 / OSL 33)
Per-change A/B from the commit messages, measured on the #20878 branch (older main). Each row was measured on the stack
of that time, so baselines differ between rows. C1/C8 = concurrency 1/8.
cu_seqlensGSM8K (400 questions, serving endpoint) stayed at 98.0-98.5%. These numbers have not yet been re-measured on this
rebased branch.
Test Coverage
test_mhc.py: split DeepGEMM tactic, lagged mode against the FMA kernels, large-M fallback tactic.glm_kpool/test_kernels.py: split top-k against one FlashMLA call (including chunks without valid rows), the LSE merge,FP8 pool scores on SM90.
glm_kpool/test_glm_kpool.py: top-k widths for both alignments.test_kimi_k3_mamba_metadata.py(host and devicecu_seqlens),test_kimi_kda_verify_parity.py(per-step andone-launch verify).
test_glm5_next_contracts.py: breakable-graph registry keys and the sparse-attention graph op.PR Checklist