Repository navigation
[None][perf] GLM-5.3-Flash FP8 performance fixes on Hopper - #20878
Closed
ruocheng-nv wants to merge 14 commits into
Closed
ruocheng-nv wants to merge 14 commits into
ruocheng-nv wants to merge 14 commits into
Conversation
The flash_mla_cpp_tllm target was built with CUDA_SEPARABLE_COMPILATION, i.e. nvcc -rdc=true. Nothing in FlashMLA needs cross-TU device linking (its own setup.py builds without -rdc), and -rdc makes ptxas spill in the attention kernels: sm90 sparse prefill (d_qk 512) 0 -> 2044 B spill stores, 296 B stack sm90 sparse fp8 decode (h64) 0 -> 220 B sm100 sparse prefill (h128) 0 -> 44 B sm100 sparse decode (h64) 0 -> 12 B H200, flash_mla_sparse_fwd with 13312 queries x 64 heads x 2176 top-k rows (d_qk 512): 17.07 ms -> 7.85 ms, outputs bit-identical. Same FlashMLA commit and nvcc; only the -rdc flag differs. Signed-off-by: Ruocheng Jia <ruochengj@nvidia.com>
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>
…pper On SM90 the CUTLASS MoE path runs FP8 block-scale expert GEMMs through the grouped DeepGEMM kernel: one persistent CTA per SM that walks every expert group. At decode batch sizes most experts are empty and each SM streams its weight tiles through a single serialized pipeline, so the GEMMs stay far from the bandwidth bound. For FP8 block-scale weights, a plain SwiGLU (optionally clamped) and at most 32 tokens per rank, CutlassFusedMoE now runs a Triton path instead: - per 1x128 FP8 activation quantization, and FP8 x FP8 GEMMs with 128x128 weight scales whose grid covers only the expert blocks the batch can occupy; - SwiGLU fused with the FP8 quantization of the FC2 input; - a single-program moe_align for up to 128 (token, expert) pairs and a top-k sum kernel. moe_align_block_size also accepts expert counts that are not a power of two. Without expert parallelism the routed expert IDs are used as-is. H200, 288 experts, top-8, hidden 4096, 512 intermediate per rank (CUDA graph), full MoE: 4 tokens 64.7 us, 1 token 28.7 us; the FC1/FC2 GEMMs on the grouped DeepGEMM path take 53/47 us and 43/36 us. GLM-5.3-Flash FP8, H200 TP4, MTP-3, 13K-token prompt, concurrency 1: decode step 14.4 -> 11.7 ms, 184 -> 225.5 output tok/s per user; GSM8K unchanged. Signed-off-by: Ruocheng Jia <ruochengj@nvidia.com>
Two of the CUTLASS MoE kernels around the FP8 block-scale grouped GEMMs were latency-bound at prefill sizes: - expandInputRowsFp8BlockScaleKernel binary-searched the expert offsets for every permuted row it moved. Each block now takes a contiguous run of permuted rows, searches once for the first row and advances the expert linearly (rows are sorted by expert). - doActivation always launched 256-thread blocks, but each thread covers eight channels, so a 512-wide intermediate slice (tensor parallelism) left three quarters of every block idle. The block is now sized to the row width in whole warps. GLM-5.3-Flash FP8, H200 TP4, 13K-token prefill (45 MoE layers, 106K expanded rows): row expansion 29.6 -> 18.9 ms, activation 13.7 -> 6.6 ms; TTFT 522 -> 507 ms at concurrency 1. 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>
Packed context rows scored their pools with the BF16 Triton kernel, about 0.84 ms per sparse layer for a 13K-token prompt. On SM100-family GPUs, quantize the pool keys per row and the queries per (row, head), fold the query scales into the head weights, and score with DeepGEMM fp8_mqa_logits. TRTLLM_GLM_INDEXER_BF16=1 keeps the BF16 scorer. Generation and verification rows are unchanged. Signed-off-by: Ruocheng Jia <ruochengj@nvidia.com>
FP8 context scoring allocated logits for every packed row against the pools of every context request, plus the gathered keys and a row copy, so its transient grew with the batch's request count and visible prefixes. Startup profiling scores a fresh prompt and never reached that growth; neither did the capacity-wide BF16 scorer once FP8 replaced it in profiling. A batch with long cached prefixes could exceed the memory left after KV-cache sizing. Context rows are now scored and selected in calls planned on the host: consecutive requests share a call while its logits, row copy and gathered keys fit 2^26 FP32 elements, each call gathers only its own requests' tables, and a request that does not fit alone is split by rows. The BF16 scorer is tiled to the same bound. KV-cache sizing reserves that bound through runtime_workspace_fixed_bytes. A fresh prompt of up to 16K tokens still scores in one call. Signed-off-by: Ruocheng Jia <ruochengj@nvidia.com>
…ound The context-score bound undercounted what a scoring call holds: - Each tile's scores stayed referenced while the next call allocated, and the last tile lived on into attention. Release them after the top-k. - The FP8 cost model missed the extra 256 columns fp8_mqa_logits pads each row with, the [rows, pools] gather temporary of a multi-table call, and the key scales and slot indices. The planner and the KV-cache reserve now share one cost function that counts them. - The reserve was min(bound, max_num_tokens * max_seq_len), which ignores how many request tables a call can gather: 64 one-row continuations of 4K prompts need about 43 MB against a 1 MiB reserve. Reserve the largest call the planner can emit for the serving capacities instead. 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
branch
from
October 23, 2026 16:23
581e7cf to
5aa96d6
Compare
Collaborator
Author
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).
Build
-rdcmade ptxas spill 2 KB per thread in theSM90 sparse prefill kernel: 17.07 -> 7.85 ms per call at T=13312.
Attention / indexer
top-k axis two blocks at a time.
(one Triton merge kernel), so small decode batches fill the GPU.
fp8_mqa_logits, with bounded per-call memory,and enable it on SM90 (
TRTLLM_GLM_INDEXER_BF16=1keeps the BF16 scorer). 1.20 -> 0.35 ms persparse layer at a 13K-token prompt. The first three of these four commits are also in [None][perf] Speed up GLM-5.3-Flash serving on Blackwell #20755;
they are included so this branch builds on main by itself.
mHC
pre-mix supported, fallback for untuned large-M buckets).
MoE (FP8 block scale)
decode-sized batches (<= 32 tokens) on Hopper.
activation blocks.
KDA
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 text serving): KDA layers join the shared mixer's breakable-graphregistry, sparse attention runs its phase split as one eager op between graph segments on the
live rows, and eager passes over a padded batch run only the live rows. Short prompts were
host-bound in eager prefill: C1 TTFT at ISL 1K 130 -> 86 ms; 4K and 13K unchanged. Buckets above
~2K tokens are not recommended (padding cost exceeds the host savings).
Performance (H200 x4, TP4, MTP3, ISL 13312 / OSL 33)
GSM8K (400 questions, serving endpoint): 98.5%.
Test Coverage
test_mhc.py(split DeepGEMM tactic, lagged path, fallback),glm_kpool/test_kernels.py(split top-k merge, FP8 pool scores incl. SM90, bounded scoring),glm_kpool/test_glm_kpool.py, KDA verify parity (per-step vs batched), Triton FP8 MoE test(added to
l0_h100.yml).glm_kpooltests on H200 with this branch: 192 passed, 6 skipped.test_glm5_next_contracts.py: breakable-graph registry keys and the sparse-attention graph op(live rows written, padded rows zeroed). GSM8K with breakable prefill graphs: 98.0%.
PR Checklist