From 904346e6dd36c26d2fd2b9da5d846226fe86a34a Mon Sep 17 00:00:00 2001 From: Ruocheng Jia Date: Thu, 22 Oct 2026 05:29:03 -0700 Subject: [PATCH 01/14] [None][perf] Build FlashMLA without relocatable device code 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 --- cpp/tensorrt_llm/flash_mla/CMakeLists.txt | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/cpp/tensorrt_llm/flash_mla/CMakeLists.txt b/cpp/tensorrt_llm/flash_mla/CMakeLists.txt index 84fff15b189e..9f552d82c0c0 100644 --- a/cpp/tensorrt_llm/flash_mla/CMakeLists.txt +++ b/cpp/tensorrt_llm/flash_mla/CMakeLists.txt @@ -197,8 +197,7 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) ) endif() -# Disable LTO before creating target (similar to DeepEP) Let CMake generate -# fatbinData for CUDA separable compilation +# Disable LTO before creating target (similar to DeepEP) set(CMAKE_INTERPROCEDURAL_OPTIMIZATION FALSE) pybind11_add_module(flash_mla_cpp_tllm ${FLASH_MLA_SOURCES}) @@ -210,8 +209,10 @@ set_target_properties( CXX_SCAN_FOR_MODULES OFF CUDA_STANDARD 20 CUDA_STANDARD_REQUIRED ON - CUDA_SEPARABLE_COMPILATION ON - CUDA_RESOLVE_DEVICE_SYMBOLS ON + # Whole-program compilation, as FlashMLA's own build does: + # relocatable device code makes ptxas spill registers in the + # attention kernels. + CUDA_SEPARABLE_COMPILATION OFF LINK_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/flash_mla_cpp_tllm.version INSTALL_RPATH "${TORCH_INSTALL_PREFIX}/lib" BUILD_WITH_INSTALL_RPATH TRUE) From 36ecd7dd05f4b41b3f90d055f3cdba997ad4c168 Mon Sep 17 00:00:00 2001 From: Ruocheng Jia Date: Thu, 22 Oct 2026 02:25:58 -0700 Subject: [PATCH 02/14] [None][perf] Use DeepGEMM TF32 for mHC layer boundaries on Hopper 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 --- tensorrt_llm/_torch/modules/mhc/mhc_cuda.py | 112 +++++++++++++++++-- tests/unittest/_torch/modules/test_mhc.py | 113 +++++++++++++++++++- 2 files changed, 215 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py b/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py index 5447d6b51afe..c37f905bb036 100644 --- a/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py +++ b/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py @@ -333,9 +333,7 @@ def get_valid_tactics( for bs in _BIGFUSE_BLOCK_SIZE_OPTIONS: tactics.append(("fma", 8, 2, bs)) - # V4 keeps its DeepGEMM tactics; the lagged (V4.1) entry boundary is pinned to the - # in-tree FMA kernel (no MMA variant exists for the pre-mapping op). - if _get_dg_fn() is not None and kwargs.get("pre_mix_ext") is None: + if _get_dg_fn() is not None: for bs in _BIGFUSE_BLOCK_SIZE_OPTIONS: tactics.append(("dg_splitk", 0, 0, bs)) tactics.append(("dg_nosplit", 0, 0, bs)) @@ -357,8 +355,6 @@ def forward(self, inputs: List[torch.Tensor], *, tactic: Any = -1, **kwargs) -> if tactic == -1: tactic = _FALLBACK_TACTIC backend, tile_n, tile_m, bigfuse_bs = tactic - if kwargs.get("pre_mix_ext") is not None and backend != "fma": - backend, tile_n, tile_m, bigfuse_bs = _FALLBACK_TACTIC num_splits = 1 if backend == "dg_splitk": @@ -825,11 +821,18 @@ def _alloc_fused_hc_scratch( _FUSED_HC_FALLBACK_TACTIC_FMA = ("fused_half_fma", 2, 1, 256, 1) -def _get_fused_hc_fallback_tactic(hidden_size: int | None = None): +def _get_fused_hc_fallback_tactic(hidden_size: int | None = None, num_tokens: int = 0): mma_ok = _fused_hc_mma_supported() if hidden_size is not None: mma_ok = mma_ok and hidden_size in _FUSED_HC_MMA_SUPPORTED_HIDDEN_SIZES - return _FUSED_HC_FALLBACK_TACTIC_MMA if mma_ok else _FUSED_HC_FALLBACK_TACTIC_FMA + if mma_ok: + return _FUSED_HC_FALLBACK_TACTIC_MMA + # Token counts above the tuned buckets (prefill chunks past + # _MHC_FINE_GRAINED_TUNING_LIMIT) land here; the FMA GEMM is several times + # slower than DeepGEMM's TF32 GEMM there. + if num_tokens > 32 and _get_dg_fn() is not None: + return ("split_dg", 0, 1 if num_tokens >= 4096 else _DG_NUM_SPLITS, 256, 1) + return _FUSED_HC_FALLBACK_TACTIC_FMA class MhcFusedHcRunner(TunableRunner): @@ -958,6 +961,13 @@ def add(tactic): for tn, ks in _FUSED_HC_HALF_FMA_TN_KS: if ks == 1: add(("fused_half_fma", tn, ks, bs, 1)) + # Without tcgen05, large M runs the boundary unfused: post_mapping, + # then the DeepGEMM TF32 GEMM + sqrsum (SM90 tensor cores) and + # big_fuse. num_k_splits carries the DeepGEMM split-K count. + if _get_dg_fn() is not None: + for splits in (1, _DG_NUM_SPLITS): + for bs in _FUSED_HC_BIGFUSE_BS: + add(("split_dg", 0, splits, bs, 1)) return tactics @@ -982,9 +992,8 @@ def forward(self, inputs, *, tactic=-1, **kwargs): pre_mix_ext = kwargs.get("pre_mix_ext") if tactic == -1: - tactic = _get_fused_hc_fallback_tactic(self.hidden_size) + tactic = _get_fused_hc_fallback_tactic(self.hidden_size, residual_prev.shape[0]) backend, tile_n, num_k_splits, bigfuse_bs, tile_m = tactic - backend_code = _FUSED_HC_BACKEND_CODE[backend] # Fused next-layer RMSNorm on layer_input. All four backends support # it (Path B/E inside mhcBigFuseKernel, Path D/F inside the all-in-one @@ -997,6 +1006,23 @@ def forward(self, inputs, *, tactic=-1, **kwargs): if norm_weight is not None and not norm_weight.is_contiguous(): norm_weight = norm_weight.contiguous() + if backend == "split_dg": + return self._forward_split_dg( + x_prev, + residual_prev, + post_mix_prev, + comb_mix_prev, + w_t_cur, + hc_scale_cur, + hc_base_cur, + num_k_splits, + bigfuse_bs, + norm_weight, + norm_eps, + pre_mix_ext, + ) + backend_code = _FUSED_HC_BACKEND_CODE[backend] + B = residual_prev.shape[0] residual_cur = torch.empty_like(residual_prev) post_mix_cur = torch.empty((B, self.n), dtype=torch.float32, device=x_prev.device) @@ -1055,6 +1081,74 @@ def forward(self, inputs, *, tactic=-1, **kwargs): return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur, pre_mix_out return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur + def _forward_split_dg( + self, + x_prev, + residual_prev, + post_mix_prev, + comb_mix_prev, + w_t_cur, + hc_scale_cur, + hc_base_cur, + num_splits, + bigfuse_bs, + norm_weight, + norm_eps, + pre_mix_ext, + ): + """post_mapping, then pre_mapping's DeepGEMM GEMM + sqrsum and big_fuse.""" + n, hidden = self.n, self.hidden_size + B = residual_prev.shape[0] + pre_mix_out = None + if pre_mix_ext is not None: + if pre_mix_ext.numel() != B * n: + # Autotuner bucket profiling (see MhcPreMappingRunner.forward). + pre_mix_ext = torch.ones((B, n), dtype=torch.float32, device=x_prev.device) + pre_mix_ext = pre_mix_ext.reshape(B, n).to(torch.float32).contiguous() + pre_mix_out = torch.empty((B, n), dtype=torch.float32, device=x_prev.device) + K = n * hidden + residual_cur = mhc_post_mapping_cuda( + residual_prev.view(B, n, hidden), + x_prev, + post_mix_prev.view(B, n), + comb_mix_prev.view(B, n, n), + n, + ) + y_acc, r_acc, num_splits = mhc_gemm_rms_dg_cuda( + residual_cur.view(B, K), w_t_cur, B, w_t_cur.shape[0], K, num_splits=num_splits + ) + post_mix_cur = torch.empty((B, n), dtype=torch.float32, device=x_prev.device) + comb_mix_cur = torch.empty((B, n * n), dtype=torch.float32, device=x_prev.device) + layer_input_cur = torch.empty_like(x_prev) + mhc_big_fuse_cuda( + y_acc, + r_acc, + residual_cur, + hc_scale_cur, + hc_base_cur, + post_mix_cur, + comb_mix_cur, + layer_input_cur, + B, + K, + hidden, + self.rms_eps, + self.hc_pre_eps, + self.hc_sinkhorn_eps, + self.hc_post_mult_value, + self.sinkhorn_repeat, + num_splits=num_splits, + block_size=bigfuse_bs, + norm_weight=norm_weight, + norm_eps=norm_eps, + pre_mix_ext=pre_mix_ext, + pre_mix_out=pre_mix_out, + ) + residual_cur = residual_cur.view(B, n, hidden) + if pre_mix_out is not None: + return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur, pre_mix_out + return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur + # Process-wide runner cache keyed on the mHC configuration. Avoids recreating # a MhcFusedHcRunner on every call. diff --git a/tests/unittest/_torch/modules/test_mhc.py b/tests/unittest/_torch/modules/test_mhc.py index 285495103a1f..fd5f451f6861 100644 --- a/tests/unittest/_torch/modules/test_mhc.py +++ b/tests/unittest/_torch/modules/test_mhc.py @@ -600,6 +600,35 @@ def test_mhc_fused_hc_single_wave_target_mma_ks(): assert ((m + 63) // 64) * ks <= sm_count * 5, (hidden_size, m, ks) +def test_mhc_fused_hc_split_dg_tactics_without_mma(monkeypatch): + """Without tcgen05, M > 32 also offers the unfused DeepGEMM boundary.""" + from tensorrt_llm._torch.autotuner import OptimizationProfile + from tensorrt_llm._torch.modules.mhc import mhc_cuda + + monkeypatch.setattr(mhc_cuda, "_fused_hc_mma_supported", lambda: False) + monkeypatch.setattr(mhc_cuda, "_get_dg_fn", lambda: object()) + runner = mhc_cuda.MhcFusedHcRunner( + n=4, + hidden_size=4096, + rms_eps=1e-6, + hc_pre_eps=1e-6, + hc_sinkhorn_eps=1e-6, + hc_post_mult_value=2.0, + sinkhorn_repeat=20, + ) + + def backends(num_tokens, **kwargs): + inputs = [torch.empty((num_tokens, 4096), dtype=torch.bfloat16)] + return {t[0] for t in runner.get_valid_tactics(inputs, OptimizationProfile(), **kwargs)} + + assert "split_dg" not in backends(32) + assert "split_dg" in backends(64) + assert "split_dg" in backends(64, pre_mix_ext=torch.empty(64, 4)) + # Token counts above the tuned buckets fall back to split_dg as well. + assert mhc_cuda._get_fused_hc_fallback_tactic(4096, 32)[0] == "fused_half_fma" + assert mhc_cuda._get_fused_hc_fallback_tactic(4096, 13312)[0] == "split_dg" + + def test_mhc_fused_hc_uneven_splits_unreachable_off_sm107(): """The uneven split counts must stay unselectable on non-SM107 devices.""" from tensorrt_llm._torch.modules.mhc.mhc_cuda import _fused_hc_target_mma_ks @@ -864,15 +893,19 @@ def make_runner_inputs(): ("fused_half_fma", 2, 1, 256, 1), pytest.param(("fused_all_mma", 0, 1, 0, 1), marks=skip_pre_blackwell), ("fused_all_fma", 2, 1, 0, 1), + ("split_dg", 0, 1, 256, 1), + ("split_dg", 0, 16, 256, 1), ], ) @pytest.mark.parametrize("hidden_size", [4096, 7168]) def test_mhc_fused_hc_realistic_scale_regression(tactic, hidden_size: int): """Real-scale mHC data catches fused_hc RMS normalization regressions.""" - from tensorrt_llm._torch.modules.mhc.mhc_cuda import MhcFusedHcRunner + from tensorrt_llm._torch.modules.mhc.mhc_cuda import MhcFusedHcRunner, _get_dg_fn if tactic[0].endswith("_mma") and not _mhc_fused_hc_mma_available(): pytest.skip("mHC fused-HC MMA kernels require SM100 and BUILD_DEEP_GEMM=ON") + if tactic[0] == "split_dg" and _get_dg_fn() is None: + pytest.skip("split_dg needs DeepGEMM's TF32 hc_prenorm GEMM") n = 16 hc_mult = 4 @@ -934,6 +967,84 @@ def test_mhc_fused_hc_realistic_scale_regression(tactic, hidden_size: int): torch.testing.assert_close(layer_input_ref, layer_input_cur, rtol=1e-2, atol=2e-2) +@pytest.mark.parametrize("num_splits", [1, 16]) +@pytest.mark.parametrize("num_tokens", [64, 1024]) +def test_mhc_split_dg_lagged_matches_fma(num_tokens: int, num_splits: int): + """With a lagged pre-mix, the DeepGEMM boundaries match the FMA ones.""" + from tensorrt_llm._torch.modules.mhc.mhc_cuda import ( + MhcFusedHcRunner, + MhcPreMappingRunner, + _get_dg_fn, + ) + + if _get_dg_fn() is None: + pytest.skip("split_dg needs DeepGEMM's TF32 hc_prenorm GEMM") + + hidden_size, hc_mult = 4096, 4 + pre_data = generate_realistic_pre_data(n=num_tokens, hc_mult=hc_mult, hidden_size=hidden_size) + consts = dict( + n=hc_mult, + hidden_size=hidden_size, + rms_eps=pre_data["rms_eps"], + hc_pre_eps=pre_data["hc_pre_eps"], + hc_sinkhorn_eps=pre_data["hc_sinkhorn_eps"], + hc_post_mult_value=pre_data["hc_post_mult_value"], + sinkhorn_repeat=pre_data["sinkhorn_repeat"], + ) + torch.random.manual_seed(23) + device = "cuda" + x_prev = torch.randn((num_tokens, hidden_size), device=device).bfloat16() + residual_prev = torch.randn((num_tokens, hc_mult, hidden_size), device=device).bfloat16() + post_mix_prev = torch.randn((num_tokens, hc_mult), dtype=torch.float32, device=device) * 0.1 + comb_mix_prev = ( + torch.randn((num_tokens, hc_mult, hc_mult), dtype=torch.float32, device=device) * 0.1 + ) + pre_mix = torch.rand((num_tokens, hc_mult), dtype=torch.float32, device=device) + w_t = pre_data["fn"].reshape(-1, hc_mult * hidden_size).contiguous() + norm_weight = (torch.rand(hidden_size, device=device) + 0.5).bfloat16() + kwargs = dict(pre_mix_ext=pre_mix, norm_weight=norm_weight, norm_eps=1e-6) + names = ("post_mix", "comb_mix", "layer_input", "pre_mix_out") + tols = ( + dict(rtol=3e-3, atol=5e-3), + dict(rtol=3e-3, atol=5e-3), + dict(rtol=1e-2, atol=2e-2), + dict(rtol=3e-3, atol=5e-3), + ) + + fused = MhcFusedHcRunner(**consts) + inputs = [ + x_prev, + residual_prev, + post_mix_prev, + comb_mix_prev, + w_t, + pre_data["hc_scale"], + pre_data["hc_base"], + ] + ref = fused(inputs=inputs, tactic=("fused_half_fma", 2, 1, 256, 1), **kwargs) + out = fused(inputs=inputs, tactic=("split_dg", 0, num_splits, 256, 1), **kwargs) + assert len(out) == len(ref) == 5 + torch.testing.assert_close(out[0], ref[0], rtol=1e-2, atol=2e-2) + for name, tol, o, r in zip(names, tols, out[1:], ref[1:]): + torch.testing.assert_close(o, r, **tol, msg=f"fused_hc {name} mismatch") + + pre_mapping = MhcPreMappingRunner(**consts) + residual = ref[0] + pre_inputs = [ + residual.view(num_tokens, hc_mult * hidden_size), + w_t, + residual, + pre_data["hc_scale"], + pre_data["hc_base"], + ] + dg_backend = "dg_splitk" if num_splits > 1 else "dg_nosplit" + ref = pre_mapping(inputs=pre_inputs, tactic=("fma", 8, 1, 256), **kwargs) + out = pre_mapping(inputs=pre_inputs, tactic=(dg_backend, 0, 0, 256), **kwargs) + assert len(out) == len(ref) == 4 + for name, tol, o, r in zip(names, tols, out, ref): + torch.testing.assert_close(o, r, **tol, msg=f"pre_mapping {name} mismatch") + + @pytest.mark.parametrize( ("num_tokens", "hidden_size", "num_k_splits"), [ From dc9ea0542e31f1c0c847375ce2b7d55aa0044f2a Mon Sep 17 00:00:00 2001 From: Ruocheng Jia Date: Wed, 21 Oct 2026 23:35:26 -0700 Subject: [PATCH 03/14] [None][fix] Align GLM-5.3-Flash sparse top-k rows for SM90 FlashMLA 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 --- .../attention/backends/sparse/glm_kpool/backend.py | 14 ++++++++++---- .../attention/backends/sparse/glm_kpool/params.py | 9 ++++----- .../attention/sparse/glm_kpool/test_glm_kpool.py | 8 +++++--- .../attention/sparse/glm_kpool/test_kernels.py | 1 + 4 files changed, 20 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py index 13937deed749..52759fce0d39 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py @@ -125,8 +125,10 @@ class GlmKpoolSparseAttention(TrtllmAttention): #: FlashMLA tiles the top-k axis in blocks of 64 (``B_TOPK``); index rows #: are padded to that multiple with ``-1`` (invalid) entries, which is - #: semantics-free by the kernel's own contract. + #: semantics-free by the kernel's own contract. Its SM90 sparse prefill + #: walks two blocks per step and requires a multiple of 128. _KERNEL_TOPK_ALIGN = 64 + _SM90_KERNEL_TOPK_ALIGN = 128 # FlashMLA instantiates 64/128 query heads. Pad TP-local heads to the next # size and discard the extra outputs; heads attend independently. @@ -205,6 +207,10 @@ def __init__( #: score scale is unchanged. self.softmax_scale = float(self.qk_nope_head_dim) ** -0.5 self._native_decode = GlmKpoolNativeDecode() + self._kernel_topk_align = ( + self._SM90_KERNEL_TOPK_ALIGN if get_sm_version() == 90 else self._KERNEL_TOPK_ALIGN + ) + self._kernel_output_width = sparse_params.kernel_output_width(self._kernel_topk_align) #: ``softmax_scale * kv_scale_quant_orig`` for the native FP8 kernel, read to #: the host once outside CUDA graph capture (the KV scale is fixed after #: construction). A float lets the kernel take it as an argument, without a @@ -434,7 +440,7 @@ def expand_selection( """Selected pools ``[N, select_k]`` -> latent row ids for the kernel. Expands each valid pool into its members, appends the always-visible - tail, pads with ``-1`` to :attr:`GlmKpoolSparseParams.kernel_output_width` + tail, pads with ``-1`` to the kernel's aligned top-k width and translates positions to latent-cache row ids in one kernel. Same row addressing as :meth:`score_pools`. """ @@ -449,7 +455,7 @@ def expand_selection( base_row=base_row, rows_per_slot=rows_per_slot, kpool=self.sparse_params.index_kpool, - out_width=self.sparse_params.kernel_output_width, + out_width=self._kernel_output_width, request_ids=request_ids, ) @@ -507,7 +513,7 @@ def _dispatch_sparse_core( q_latent[start:stop], selected, local_indices ) return output - pad = (-topk_rows.shape[-1]) % self._KERNEL_TOPK_ALIGN + pad = (-topk_rows.shape[-1]) % self._kernel_topk_align if pad: topk_rows = torch.nn.functional.pad(topk_rows, (0, pad), value=INDEX_SENTINEL) local_heads = q_latent.shape[1] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/params.py b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/params.py index b43a9881d4a1..181b8f6bb6da 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/params.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/params.py @@ -88,10 +88,9 @@ def select_k(self) -> int: """Number of pools a query selects.""" return self.index_topk // self.index_kpool - @property - def kernel_output_width(self) -> int: - """``output_width`` padded to the FlashMLA top-k tile (64).""" - return -(-self.output_width // 64) * 64 + def kernel_output_width(self, topk_align: int) -> int: + """``output_width`` padded to the sparse kernel's top-k alignment.""" + return -(-self.output_width // topk_align) * topk_align @property def output_width(self) -> int: @@ -103,7 +102,7 @@ def output_width(self) -> int: class GlmKpoolBackendForwardArgs(SparseBackendForwardArgs): """Global latent-cache row selection produced by expand_selection. - topk_rows is int32 [tokens, kernel_output_width], with -1 marking invalid + topk_rows is int32 [tokens, kernel_output_width(align)], with -1 marking invalid entries. Request-local topk_indices from the base carrier are unsupported. """ diff --git a/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py b/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py index e989db333636..c1bbb6d5143b 100644 --- a/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py +++ b/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py @@ -455,12 +455,14 @@ def test_derived_geometry_follows_the_checkpoint_defaults(): assert params.select_k == 2048 // 4 # topk positions plus the always-visible incomplete tail (kpool - 1 rows). assert params.output_width == 2048 + 3 - # Padded to the FlashMLA top-k tile. - assert params.kernel_output_width == 2112 + # Padded to the FlashMLA top-k tile (64), or to two tiles for its SM90 prefill. + assert params.kernel_output_width(64) == 2112 + assert params.kernel_output_width(128) == 2176 assert params.packed_state_dim == 2 * 128 assert params.cache_row_dim == 3 * 128 small = GlmKpoolSparseParams(index_topk=64, index_kpool=8) - assert (small.select_k, small.output_width, small.kernel_output_width) == (8, 71, 128) + assert (small.select_k, small.output_width) == (8, 71) + assert small.kernel_output_width(64) == small.kernel_output_width(128) == 128 def test_paged_slot_indices_returns_page_and_offset_pairs(): diff --git a/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py b/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py index de5148df592b..0dcb75dd4ae4 100644 --- a/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py +++ b/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py @@ -382,6 +382,7 @@ def test_fp8_sparse_core_matches_dequantized_cache_and_replays_graph(heads): backend = object.__new__(GlmKpoolSparseAttention) backend.kv_lora_rank = 512 backend.softmax_scale = 256**-0.5 + backend._kernel_topk_align = 128 if get_sm_version() == 90 else 64 backend.kv_scale_quant_orig = torch.tensor([0.5], device="cuda") backend._FP8_QUERY_CHUNK_SIZE = 2 # Exercise multiple chunks and a short final chunk. generator = torch.Generator(device="cuda").manual_seed(71) From e0c7d60dac7fec0420500ee54edd484725f425e7 Mon Sep 17 00:00:00 2001 From: Ruocheng Jia Date: Thu, 22 Oct 2026 19:14:22 -0700 Subject: [PATCH 04/14] [None][perf] Split the sparse top-k axis for decode-sized calls on SM90 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 --- .../backends/sparse/glm_kpool/backend.py | 40 ++++++++++++ .../backends/sparse/glm_kpool/kernels.py | 63 +++++++++++++++++++ .../sparse/glm_kpool/test_kernels.py | 55 ++++++++++++++++ 3 files changed, 158 insertions(+) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py index 52759fce0d39..97225eb61c2a 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py @@ -47,6 +47,7 @@ kpool_expand, kpool_score, kpool_update, + merge_split_attention, ) from .metadata import Glm5NextMamba2Metadata from .native_decode import GlmKpoolNativeDecode @@ -129,6 +130,8 @@ class GlmKpoolSparseAttention(TrtllmAttention): #: walks two blocks per step and requires a multiple of 128. _KERNEL_TOPK_ALIGN = 64 _SM90_KERNEL_TOPK_ALIGN = 128 + #: Calls with at most this many query tokens split the top-k axis on SM90. + _SM90_SPLIT_TOPK_MAX_TOKENS = 16 # FlashMLA instantiates 64/128 query heads. Pad TP-local heads to the next # size and discard the extra outputs; heads attend independently. @@ -211,6 +214,9 @@ def __init__( self._SM90_KERNEL_TOPK_ALIGN if get_sm_version() == 90 else self._KERNEL_TOPK_ALIGN ) self._kernel_output_width = sparse_params.kernel_output_width(self._kernel_topk_align) + #: The SM90 sparse kernel runs one CTA per query token, so a decode-sized + #: call occupies a handful of SMs; split its top-k axis across CTAs. + self._split_topk = get_sm_version() == 90 #: ``softmax_scale * kv_scale_quant_orig`` for the native FP8 kernel, read to #: the host once outside CUDA graph capture (the KV scale is fixed after #: construction). A float lets the kernel take it as an argument, without a @@ -531,6 +537,15 @@ def _dispatch_sparse_core( q_latent = q_padded if flash_mla_sparse_fwd is None: raise RuntimeError("glm_kpool sparse attention requires FlashMLA in this build") + num_tokens, width = q_latent.shape[0], topk_rows.shape[-1] + if ( + self._split_topk + and num_tokens <= self._SM90_SPLIT_TOPK_MAX_TOKENS + and width > self._kernel_topk_align + ): + return self._split_topk_sparse( + q_latent, kv_rows, topk_rows, width // self._kernel_topk_align, local_heads + ) out, _, _ = flash_mla_sparse_fwd( q_latent, kv_rows, @@ -544,6 +559,31 @@ def _dispatch_sparse_core( out = out[:, :local_heads, :] return out + def _split_topk_sparse( + self, + q_latent: torch.Tensor, + kv_rows: torch.Tensor, + topk_rows: torch.Tensor, + splits: int, + local_heads: int, + ) -> torch.Tensor: + """Sparse attention over ``splits`` equal top-k chunks, merged by LSE. + + Each (token, chunk) pair is one kernel row, so the call spreads over + ``splits`` times as many CTAs; returns ``[T, local_heads, kv_lora]``. + """ + num_tokens, heads, dim = q_latent.shape + chunk = topk_rows.shape[-1] // splits + q_rep = q_latent.unsqueeze(1).expand(num_tokens, splits, heads, dim) + out, _, lse = flash_mla_sparse_fwd( + q_rep.reshape(num_tokens * splits, heads, dim), + kv_rows, + topk_rows.view(num_tokens * splits, 1, chunk), + self.softmax_scale, + self.kv_lora_rank, + ) + return merge_split_attention(out, lse, num_tokens, splits, local_heads) + def _finalize_output( self, out_latent: torch.Tensor, output: torch.Tensor | None ) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/kernels.py b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/kernels.py index 02297052194f..2a6c56c21283 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/kernels.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/kernels.py @@ -683,3 +683,66 @@ def compact_sparse_rows(indices: torch.Tensor) -> tuple[torch.Tensor, torch.Tens triton.next_power_of_2(width), ) return packed, lengths, nonempty + + +@triton.jit +def _merge_split_attention_kernel( + OUT, + LSE, + DST, + stride_out_row, + stride_out_head, + stride_lse_row, + stride_dst_tok, + stride_dst_head, + SPLITS: tl.constexpr, + DIM: tl.constexpr, +): + tok = tl.program_id(0) + head = tl.program_id(1) + cols = tl.arange(0, DIM) + max_lse = float("-inf") + for s in range(SPLITS): + lse = tl.load(LSE + (tok * SPLITS + s) * stride_lse_row + head) + # A chunk without valid rows reports +inf (and a zero output). + lse = tl.where(lse == float("inf"), float("-inf"), lse) + max_lse = tl.maximum(max_lse, lse) + total = 0.0 + acc = tl.zeros([DIM], dtype=tl.float32) + for s in range(SPLITS): + row = tok * SPLITS + s + lse = tl.load(LSE + row * stride_lse_row + head) + lse = tl.where(lse == float("inf"), float("-inf"), lse) + w = tl.where(max_lse == float("-inf"), 0.0, tl.exp(lse - max_lse)) + part = tl.load(OUT + row * stride_out_row + head * stride_out_head + cols).to(tl.float32) + acc += w * part + total += w + acc = tl.where(total > 0.0, acc / total, 0.0) + tl.store( + DST + tok * stride_dst_tok + head * stride_dst_head + cols, acc.to(DST.dtype.element_ty) + ) + + +def merge_split_attention( + out: torch.Tensor, lse: torch.Tensor, num_tokens: int, splits: int, heads: int +) -> torch.Tensor: + """Merge per-chunk attention ``[tokens * splits, H, D]`` by LSE into ``[tokens, heads, D]``. + + ``lse`` is ``[tokens * splits, H]``; chunks without valid rows (LSE ``+inf``) + contribute nothing. Only the first ``heads`` query heads are merged. + """ + dst = torch.empty((num_tokens, heads, out.shape[-1]), dtype=out.dtype, device=out.device) + if num_tokens: + _merge_split_attention_kernel[(num_tokens, heads)]( + out, + lse, + dst, + out.stride(0), + out.stride(1), + lse.stride(0), + dst.stride(0), + dst.stride(1), + SPLITS=splits, + DIM=out.shape[-1], + ) + return dst diff --git a/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py b/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py index 0dcb75dd4ae4..b13dab5487ce 100644 --- a/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py +++ b/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py @@ -23,6 +23,7 @@ kpool_expand, kpool_score, kpool_update, + merge_split_attention, ) from tensorrt_llm._utils import get_sm_version @@ -383,6 +384,7 @@ def test_fp8_sparse_core_matches_dequantized_cache_and_replays_graph(heads): backend.kv_lora_rank = 512 backend.softmax_scale = 256**-0.5 backend._kernel_topk_align = 128 if get_sm_version() == 90 else 64 + backend._split_topk = False backend.kv_scale_quant_orig = torch.tensor([0.5], device="cuda") backend._FP8_QUERY_CHUNK_SIZE = 2 # Exercise multiple chunks and a short final chunk. generator = torch.Generator(device="cuda").manual_seed(71) @@ -415,6 +417,59 @@ def reference(): torch.testing.assert_close(captured, reference(), atol=2e-2, rtol=2e-2) +@pytest.mark.skipif(not HAS_FLASH_MLA, reason="FlashMLA not available") +@pytest.mark.skipif(get_sm_version() < 90, reason="FlashMLA requires SM90 (Hopper) or later") +@pytest.mark.parametrize("num_tokens", [1, 4, 16]) +def test_split_topk_sparse_core_matches_single_call(num_tokens): + """Decode-sized calls split over top-k chunks match the single-call kernel.""" + from tensorrt_llm._torch.attention.backends.sparse.glm_kpool import GlmKpoolSparseAttention + + def make_backend(split: bool) -> GlmKpoolSparseAttention: + backend = object.__new__(GlmKpoolSparseAttention) + backend.kv_lora_rank = 512 + backend.softmax_scale = 256**-0.5 + backend._kernel_topk_align = 128 if get_sm_version() == 90 else 64 + backend._split_topk = split + return backend + + generator = torch.Generator(device="cuda").manual_seed(5) + kv_rows = (torch.randn(4096, 1, 512, device="cuda", generator=generator) * 0.5).to( + torch.bfloat16 + ) + q = (torch.randn(num_tokens, 16, 512, device="cuda", generator=generator) * 0.5).to( + torch.bfloat16 + ) + indices = torch.randint( + 0, kv_rows.shape[0], (num_tokens, 2112), device="cuda", generator=generator + ).int() + indices[:, -61:] = -1 # Padding tail, as in the expanded k-pool rows. + indices[0, : 2 * 128] = -1 # Whole chunks without a valid row. + + expected = make_backend(split=False)._dispatch_sparse_core(q, kv_rows, indices) + actual = make_backend(split=True)._dispatch_sparse_core(q, kv_rows, indices) + assert actual.shape == expected.shape == (num_tokens, 16, 512) + torch.testing.assert_close(actual, expected, atol=2e-3, rtol=2e-2) + + +def test_merge_split_attention_matches_reference(): + generator = torch.Generator(device="cuda").manual_seed(9) + num_tokens, splits, heads, kernel_heads, dim = 3, 5, 16, 64, 512 + out = torch.randn( + num_tokens * splits, kernel_heads, dim, device="cuda", generator=generator + ).to(torch.bfloat16) + lse = torch.randn(num_tokens * splits, kernel_heads, device="cuda", generator=generator) * 4 + lse[1] = float("inf") # A chunk without valid rows. + + actual = merge_split_attention(out, lse, num_tokens, splits, heads) + + lse_ref = lse.view(num_tokens, splits, kernel_heads)[..., :heads] + lse_ref = torch.where(torch.isinf(lse_ref), torch.full_like(lse_ref, float("-inf")), lse_ref) + weights = torch.softmax(lse_ref, dim=1) + parts = out.view(num_tokens, splits, kernel_heads, dim)[:, :, :heads].float() + expected = (weights[..., None] * parts).sum(dim=1) + torch.testing.assert_close(actual.float(), expected, atol=2e-2, rtol=2e-2) + + @pytest.mark.parametrize("context", [False, True]) def test_fp8_cache_storage_helpers_preserve_coalesced_pages(context): """Write FP8 payloads without disturbing coalesced pages or indexer pool keys.""" From 79bd1fd27df015a327b941aebf09c7ae22320fe7 Mon Sep 17 00:00:00 2001 From: Ruocheng Jia Date: Thu, 22 Oct 2026 08:57:59 -0700 Subject: [PATCH 05/14] [None][perf] Triton FP8-activation MoE for decode-sized batches on Hopper 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 --- .../_torch/moe/fused_moe/fused_moe_cutlass.py | 72 ++- .../fused_moe_triton_fp8_block_scale.py | 460 +++++++++++++++++- .../test_lists/test-db/l0_h100.yml | 1 + .../fused_moe/test_triton_fp8_block_scale.py | 114 ++++- 4 files changed, 635 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py index 8d78c8c57130..b6f3b0843ab3 100755 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py @@ -28,7 +28,7 @@ from ...utils import (ActivationType, AuxStreamType, EventType, Fp4QuantizedTensor) from .activation import (DEFAULT_MOE_ACTIVATION, ActivationParamShape, - MoEActivation, MoEActivationSupport) + MoEActivation, MoEActivationSupport, SwigluActivation) from .impl_base import MoEImplBase, apply_moe_impl_construction_state from .impl_contract import (MoEDeployment, MoEEligibility, MoEInputRequirement, MoEProblem, MoERejectReason, MoERunContext, @@ -65,6 +65,11 @@ def raise_moe_lora_multichunk_unsupported(num_chunks: int) -> None: f"increase `moe_max_num_tokens` so the MoE runs in a single chunk.") +# Largest batch (tokens on this rank) that runs the SM90 FP8-activation Triton +# MoE instead of the CUTLASS / DeepGEMM grouped GEMM. +_SM90_TRITON_FP8_MOE_MAX_TOKENS = 32 + + class CutlassFusedMoE(MoEImplBase): """ Fused Mixture of Experts (MoE) Layer with performance tuning. @@ -973,6 +978,32 @@ def _tuner_shapes( tuner_num_tokens = ctx.x.shape[0] * self.mapping.tp_size return tuner_num_tokens, self.routing_method.top_k + def _triton_fp8_act_swiglu( + self, + x: torch.Tensor, + x_sf: Optional[torch.Tensor], + lora_params, + moe_output: Optional[torch.Tensor], + ) -> Optional[Tuple[Optional[float], bool]]: + """SwiGLU clamp settings when this call can take the SM90 FP8-activation Triton path. + + Returns ``(limit, clamp_after_silu)``, or None to keep the CUTLASS path: + FP8 block-scale weights, a plain SwiGLU with at most a scalar clamp, BF16 + input, a decode-sized batch, and no LoRA or caller-provided output buffer. + """ + if (not self.has_deepseek_fp8_block_scales + or self.activation_type != ActivationType.Swiglu + or not isinstance(self.activation, SwigluActivation) + or x_sf is not None or x.dtype != torch.bfloat16 + or x.shape[0] > _SM90_TRITON_FP8_MOE_MAX_TOKENS + or lora_params is not None or moe_output is not None): + return None + limit = self.activation.clamp + if isinstance(limit, torch.Tensor): + return None + return (None if limit is None else float(limit), + self.activation.clamp_after_silu) + def run_moe( self, ctx: MoERunContext, @@ -1016,9 +1047,17 @@ def run_moe( # SM120 + FP8 block scales: use Triton kernel (CUTLASS TMA fails on SM120 # for large token counts due to cuTensorMapEncodeTiled limitations). - if self.has_deepseek_fp8_block_scales and get_sm_version() == 120: - from .fused_moe_triton_fp8_block_scale import \ - run_triton_fp8_block_scale_moe + # SM90 + FP8 block scales, decode-sized batches: the Triton FP8-activation + # kernels launch only the expert blocks the batch routes to, where the + # grouped DeepGEMM kernel walks every expert on one persistent CTA per SM. + sm_version = get_sm_version() + triton_fp8_act_swiglu = (self._triton_fp8_act_swiglu( + x, x_sf, lora_params, moe_output) if sm_version == 90 else None) + if self.has_deepseek_fp8_block_scales and ( + sm_version == 120 or triton_fp8_act_swiglu is not None): + from .fused_moe_triton_fp8_block_scale import ( + run_triton_fp8_block_scale_moe, + run_triton_fp8_block_scale_moe_fp8_act) # forward_chunk sets token_final_scales=None when # apply_router_weight_on_input=True (weights already folded into x); @@ -1032,7 +1071,11 @@ def run_moe( # (0 .. expert_size_per_partition-1), so remap and zero-scale any # non-local token-expert pairs to suppress their contribution. local_n = self.expert_size_per_partition - if enable_alltoall: + if self.ep_size == 1: + # Every routed expert is local: the IDs and weights need no remap. + local_ids = token_selected_experts + is_local = None + elif enable_alltoall: # After alltoall dispatch, IDs are already local; padding = local_n local_ids = token_selected_experts.clamp(0, local_n - 1) is_local = token_selected_experts < local_n @@ -1042,8 +1085,23 @@ def run_moe( 0, local_n - 1) is_local = ((token_selected_experts >= slot_start) & (token_selected_experts < slot_start + local_n)) - local_scales = token_final_scales * is_local.to( - token_final_scales.dtype) + local_scales = (token_final_scales + if is_local is None else token_final_scales * + is_local.to(token_final_scales.dtype)) + if triton_fp8_act_swiglu is not None: + swiglu_limit, swiglu_clamp_after_silu = triton_fp8_act_swiglu + return run_triton_fp8_block_scale_moe_fp8_act( + x, + local_ids, + local_scales, + self.w3_w1_weight, + self.quant_scales.fc_weight_scales, + self.w2_weight, + self.quant_scales.proj_weight_scales, + swiglu_limit=swiglu_limit, + swiglu_clamp_after_silu=swiglu_clamp_after_silu, + output_dtype=output_dtype, + ) result = run_triton_fp8_block_scale_moe( x, local_ids, diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton_fp8_block_scale.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton_fp8_block_scale.py index 7d837e9576e0..003f141e721b 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton_fp8_block_scale.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton_fp8_block_scale.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright 2026 NVIDIA Corporation. All rights reserved. """ -Triton FP8 block-scale MoE forward pass for SM120. +Triton FP8 block-scale MoE forward passes: BF16 activations for SM120, FP8 +activations for decode-sized batches on SM90. """ import torch @@ -40,15 +41,17 @@ def _moe_prefix_kernel( num_post_pad_ptr, BLOCK_SIZE: tl.constexpr, NUM_EXPERTS: tl.constexpr, + NUM_EXPERTS_POW2: tl.constexpr, ): """Compute exclusive offsets from padded expert counts.""" - e = tl.arange(0, NUM_EXPERTS) - counts = tl.load(counts_ptr + e).to(tl.int32) + e = tl.arange(0, NUM_EXPERTS_POW2) + mask = e < NUM_EXPERTS + counts = tl.load(counts_ptr + e, mask=mask, other=0).to(tl.int32) padded = ((counts + BLOCK_SIZE - 1) // BLOCK_SIZE) * BLOCK_SIZE # Inclusive scan; subtract self for exclusive prefix. inc = tl.associative_scan(padded, 0, _int32_add) out_off = inc - padded - tl.store(out_off_ptr + e, out_off) + tl.store(out_off_ptr + e, out_off, mask=mask) tl.store(num_post_pad_ptr, tl.sum(padded).to(tl.int32)) @@ -139,6 +142,7 @@ def moe_align_block_size( num_post_pad, BLOCK_SIZE=block_size, NUM_EXPERTS=num_experts, + NUM_EXPERTS_POW2=triton.next_power_of_2(num_experts), ) _moe_expert_ids_kernel[num_experts,]( @@ -601,3 +605,451 @@ def run_triton_fp8_block_scale_moe( if output_dtype is not None and output.dtype != output_dtype: output = output.to(output_dtype) return output + + +# --------------------------------------------------------------------------- +# FP8-activation path for decode-sized batches (SM90) +# --------------------------------------------------------------------------- +_FP8_GROUP_SIZE = 128 +_FP8_E4M3_MAX = 448.0 +# Largest (token, expert) pair count aligned by the single-program kernel. +_SMALL_ALIGN_MAX_PAIRS = 128 + + +@triton.jit +def _per_token_group_quant_fp8_kernel( + x_ptr, + q_ptr, + s_ptr, + stride_xm, + stride_qm, + stride_sm, + GROUP: tl.constexpr, + FP8_MAX: tl.constexpr, +): + """Quantize one 1 x GROUP slice of a row to FP8 with its own scale.""" + row = tl.program_id(0).to(tl.int64) + group = tl.program_id(1) + offs = group * GROUP + tl.arange(0, GROUP) + x = tl.load(x_ptr + row * stride_xm + offs).to(tl.float32) + scale = tl.maximum(tl.max(tl.abs(x)), 1e-10) / FP8_MAX + tl.store(q_ptr + row * stride_qm + offs, (x / scale).to(tl.float8e4nv)) + tl.store(s_ptr + row * stride_sm + group, scale) + + +@triton.jit +def _swiglu_quant_fp8_kernel( + ic1_ptr, + q_ptr, + s_ptr, + INTER, + stride_im, + stride_qm, + stride_sm, + limit, + HAS_LIMIT: tl.constexpr, + CLAMP_AFTER_SILU: tl.constexpr, + GROUP: tl.constexpr, + FP8_MAX: tl.constexpr, +): + """SwiGLU over one 1 x GROUP slice of an (up | gate) FC1 row, quantized to FP8. + + With a limit, ``up`` is clamped to [-limit, limit] and ``gate`` to at most + ``limit`` before SiLU (or the SiLU output, with CLAMP_AFTER_SILU), as the + CUTLASS MoE activation does. + """ + row = tl.program_id(0).to(tl.int64) + group = tl.program_id(1) + offs = group * GROUP + tl.arange(0, GROUP) + up = tl.load(ic1_ptr + row * stride_im + offs).to(tl.float32) + gate = tl.load(ic1_ptr + row * stride_im + INTER + offs).to(tl.float32) + if HAS_LIMIT: + up = tl.minimum(tl.maximum(up, -limit), limit) + if CLAMP_AFTER_SILU: + act = tl.minimum(gate * tl.sigmoid(gate), limit) + else: + gate = tl.minimum(gate, limit) + act = gate * tl.sigmoid(gate) + else: + act = gate * tl.sigmoid(gate) + y = act * up + scale = tl.maximum(tl.max(tl.abs(y)), 1e-10) / FP8_MAX + tl.store(q_ptr + row * stride_qm + offs, (y / scale).to(tl.float8e4nv)) + tl.store(s_ptr + row * stride_sm + group, scale) + + +@triton.jit +def _fused_moe_fp8act_fp8w_kernel( + a_ptr, + a_scale_ptr, + b_ptr, + b_scale_ptr, + c_ptr, + topk_weights_ptr, + sorted_token_ids_ptr, + expert_ids_ptr, + num_tokens_post_padded_ptr, + N, + K, + num_valid_tokens, + stride_am, + stride_ak, + stride_asm, + stride_ask, + stride_be, + stride_bn, + stride_bk, + stride_bse, + stride_bsn, + stride_bsk, + stride_cm, + stride_cn, + group_n: tl.constexpr, + group_k: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + top_k: tl.constexpr, + MUL_ROUTED_WEIGHT: tl.constexpr, +): + """Fused MoE GEMM with FP8 activations (1 x 128 scales) and FP8 weights (128 x 128 scales). + + One program per (expert block of BLOCK_SIZE_M sorted rows, BLOCK_SIZE_N + output columns); C row ``r`` holds the (token, expert) slot ``r``. + C[m, n] = sum_k dot(A_fp8[m, k], B_fp8[k, n]) * A_scale[m, k] * B_scale[n, k] + """ + pid = tl.program_id(0) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + pid_m = pid // num_pid_n + pid_n = pid % num_pid_n + if pid_m * BLOCK_SIZE_M >= tl.load(num_tokens_post_padded_ptr): + return + + offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + offs_token = tl.load(sorted_token_ids_ptr + offs_token_id).to(tl.int64) + token_mask = offs_token < num_valid_tokens + off_expert = tl.load(expert_ids_ptr + pid_m).to(tl.int64) + + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_row = offs_token // top_k + a_ptrs = a_ptr + a_row[:, None] * stride_am + offs_k[None, :] * stride_ak + b_ptrs = ( + b_ptr + off_expert * stride_be + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn + ) + a_scale_ptrs = a_scale_ptr + a_row * stride_asm + b_scale_ptrs = b_scale_ptr + off_expert * stride_bse + (offs_bn // group_n) * stride_bsn + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + a = tl.load(a_ptrs, mask=token_mask[:, None], other=0.0) + b = tl.load(b_ptrs) + offs_ks = (k * BLOCK_SIZE_K) // group_k + a_scale = tl.load(a_scale_ptrs + offs_ks * stride_ask, mask=token_mask, other=0.0) + b_scale = tl.load(b_scale_ptrs + offs_ks * stride_bsk) + accumulator += tl.dot(a, b) * a_scale[:, None] * b_scale[None, :] + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + if MUL_ROUTED_WEIGHT: + moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0.0) + accumulator *= moe_weight[:, None] + + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :] + c_mask = token_mask[:, None] & (offs_cn[None, :] < N) + tl.store(c_ptrs, accumulator.to(tl.bfloat16), mask=c_mask) + + +@triton.jit +def _moe_align_small_kernel( + flat_ids_ptr, + sorted_token_ids_ptr, + expert_ids_ptr, + num_post_pad_ptr, + expert_offsets_ptr, + num_pairs, + num_slots, + NUM_PAIRS_POW2: tl.constexpr, + NUM_SLOTS_POW2: tl.constexpr, + NUM_EXPERTS_POW2: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """moe_align_block_size in one program for small batches. + + A pair's slot is its expert's padded offset (staged in ``expert_offsets``) + plus the number of earlier pairs routed to the same expert, so the pairwise + comparison replaces the atomic scatter of the multi-kernel version. + """ + i = tl.arange(0, NUM_PAIRS_POW2) + valid = i < num_pairs + # Padding pairs go to the last bin, which no real expert uses. + ids = tl.load(flat_ids_ptr + i, mask=valid, other=NUM_EXPERTS_POW2 - 1).to(tl.int32) + e = tl.arange(0, NUM_EXPERTS_POW2) + counts = tl.histogram(ids, NUM_EXPERTS_POW2) + padded = ((counts + BLOCK_SIZE - 1) // BLOCK_SIZE) * BLOCK_SIZE + padded = tl.where(e == NUM_EXPERTS_POW2 - 1, 0, padded) + tl.store(expert_offsets_ptr + e, tl.cumsum(padded, axis=0) - padded) + tl.debug_barrier() + pair_offset = tl.load(expert_offsets_ptr + ids) + earlier = (ids[None, :] == ids[:, None]) & (i[None, :] < i[:, None]) & valid[None, :] + slot = pair_offset + tl.sum(earlier.to(tl.int32), axis=1) + + pad = tl.arange(0, NUM_SLOTS_POW2) + tl.store( + sorted_token_ids_ptr + pad, + tl.full([NUM_SLOTS_POW2], 0, tl.int32) + num_pairs, + mask=pad < num_slots, + ) + tl.debug_barrier() + tl.store(sorted_token_ids_ptr + slot, i, mask=valid) + tl.store(expert_ids_ptr + slot // BLOCK_SIZE, ids, mask=valid) + tl.store(num_post_pad_ptr, tl.sum(padded)) + + +def _moe_align_small( + topk_ids: torch.Tensor, block_size: int, num_experts: int, num_m_blocks: int +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Same outputs as :func:`moe_align_block_size` for the first ``num_m_blocks`` blocks.""" + device = topk_ids.device + flat_ids = topk_ids.reshape(-1) + num_pairs = flat_ids.shape[0] + num_slots = num_m_blocks * block_size + sorted_token_ids = torch.empty((num_slots,), dtype=torch.int32, device=device) + expert_ids = torch.empty((num_m_blocks,), dtype=torch.int32, device=device) + num_post_pad = torch.empty(1, dtype=torch.int32, device=device) + num_experts_pow2 = triton.next_power_of_2(num_experts + 1) + expert_offsets = torch.empty((num_experts_pow2,), dtype=torch.int32, device=device) + _moe_align_small_kernel[(1,)]( + flat_ids, + sorted_token_ids, + expert_ids, + num_post_pad, + expert_offsets, + num_pairs, + num_slots, + NUM_PAIRS_POW2=triton.next_power_of_2(max(num_pairs, 2)), + NUM_SLOTS_POW2=triton.next_power_of_2(max(num_slots, 2)), + NUM_EXPERTS_POW2=num_experts_pow2, + BLOCK_SIZE=block_size, + num_warps=4, + ) + return sorted_token_ids, expert_ids, num_post_pad + + +@triton.jit +def _moe_topk_sum_kernel( + in_ptr, + out_ptr, + hidden, + TOP_K: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """out[t, :] = sum_k in[t * TOP_K + k, :] in FP32.""" + t = tl.program_id(0).to(tl.int64) + offs = tl.program_id(1) * BLOCK_H + tl.arange(0, BLOCK_H) + mask = offs < hidden + acc = tl.zeros([BLOCK_H], dtype=tl.float32) + for k in tl.static_range(TOP_K): + acc += tl.load(in_ptr + (t * TOP_K + k) * hidden + offs, mask=mask, other=0.0).to( + tl.float32 + ) + tl.store(out_ptr + t * hidden + offs, acc.to(out_ptr.dtype.element_ty), mask=mask) + + +def _per_token_group_quant_fp8(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """(M, K) BF16 -> (M, K) FP8 and (M, K // 128) FP32 scales.""" + M, K = x.shape + assert K % _FP8_GROUP_SIZE == 0, f"K={K} must be a multiple of {_FP8_GROUP_SIZE}" + q = torch.empty((M, K), dtype=torch.float8_e4m3fn, device=x.device) + scales = torch.empty((M, K // _FP8_GROUP_SIZE), dtype=torch.float32, device=x.device) + if M > 0: + _per_token_group_quant_fp8_kernel[(M, K // _FP8_GROUP_SIZE)]( + x, + q, + scales, + x.stride(0), + q.stride(0), + scales.stride(0), + GROUP=_FP8_GROUP_SIZE, + FP8_MAX=_FP8_E4M3_MAX, + ) + return q, scales + + +def _invoke_fp8act_fp8w_moe_kernel( + a: torch.Tensor, + a_scale: torch.Tensor, + b: torch.Tensor, + b_scale: torch.Tensor, + c: torch.Tensor, + topk_weights: torch.Tensor | None, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + num_m_blocks: int, + num_valid_tokens: int, + top_k: int, +) -> None: + """Launch `_fused_moe_fp8act_fp8w_kernel` over at most ``num_m_blocks`` expert blocks. + + b: (E, N, K) FP8, b_scale: (E, N // 128, K // 128) FP32, c: (num_valid_tokens, N) BF16. + """ + N, K = b.shape[1], b.shape[2] + # Measured on H200 at GLM-5.3-Flash TP4 shapes: narrower N tiles feed more SMs when few + # expert blocks stream a long K. + if K >= 2048 and num_m_blocks <= 32: + block_n, num_warps, num_stages = 64, 4, 4 + else: + block_n, num_warps, num_stages = 128, 4, 3 + grid = (num_m_blocks * triton.cdiv(N, block_n),) + _fused_moe_fp8act_fp8w_kernel[grid]( + a, + a_scale, + b, + b_scale, + c, + topk_weights, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + N, + K, + num_valid_tokens, + a.stride(0), + a.stride(1), + a_scale.stride(0), + a_scale.stride(1), + b.stride(0), + b.stride(1), + b.stride(2), + b_scale.stride(0), + b_scale.stride(1), + b_scale.stride(2), + c.stride(0), + c.stride(1), + group_n=_BLOCK_SHAPE[1], + group_k=_BLOCK_SHAPE[0], + BLOCK_SIZE_M=_BLOCK_SIZE_M, + BLOCK_SIZE_N=block_n, + BLOCK_SIZE_K=_BLOCK_SHAPE[0], + top_k=top_k, + MUL_ROUTED_WEIGHT=topk_weights is not None, + num_warps=num_warps, + num_stages=num_stages, + ) + + +def run_triton_fp8_block_scale_moe_fp8_act( + x: torch.Tensor, + token_selected_experts: torch.Tensor, + token_final_scales: torch.Tensor, + w3_w1: torch.Tensor, + w3_w1_scales: torch.Tensor, + w2: torch.Tensor, + w2_scales: torch.Tensor, + swiglu_limit: float | None = None, + swiglu_clamp_after_silu: bool = False, + output_dtype: torch.dtype | None = None, +) -> torch.Tensor: + """SwiGLU MoE forward with FP8 activations for decode-sized batches. + + Same weight layout as :func:`run_triton_fp8_block_scale_moe` + (DeepSeekFP8BlockScalesFusedMoEMethod without E8M0 resmoothing). Activations + are quantized per 1 x 128 group, and each GEMM launches only as many expert + blocks as the batch can occupy, so the cost follows the routed experts + rather than the expert count. + + Args: + x: (T, H) BF16 input activations. + token_selected_experts: (T, K) expert IDs local to ``w3_w1``. + token_final_scales: (T, K) F32 routing weights. + w3_w1: (E, 2I, H) FP8 up+gate projection weights. + w3_w1_scales: (E, 2I // 128, H // 128) F32 block scales. + w2: (E, H, I) FP8 down-projection weights. + w2_scales: (E, H // 128, I // 128) F32 block scales. + swiglu_limit: Optional SwiGLU clamp. + swiglu_clamp_after_silu: Clamp the SiLU output instead of the gate. + output_dtype: Optional output dtype. + + Returns: + (T, H) output tensor. + """ + num_tokens, hidden = x.shape + num_experts, gate_up_size, _ = w3_w1.shape + intermediate = gate_up_size // 2 + top_k = token_selected_experts.shape[1] + num_pairs = num_tokens * top_k + device = x.device + + # Every expert block holds at least one (token, expert) pair. + num_m_blocks = min(num_pairs, num_pairs // _BLOCK_SIZE_M + num_experts) + if num_pairs <= _SMALL_ALIGN_MAX_PAIRS: + sorted_token_ids, expert_ids, num_tokens_post_padded = _moe_align_small( + token_selected_experts, _BLOCK_SIZE_M, num_experts, num_m_blocks + ) + else: + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( + token_selected_experts, _BLOCK_SIZE_M, num_experts + ) + + x_q, x_scales = _per_token_group_quant_fp8(x.contiguous()) + ic1 = torch.empty((num_pairs, gate_up_size), dtype=torch.bfloat16, device=device) + _invoke_fp8act_fp8w_moe_kernel( + x_q, + x_scales, + w3_w1, + w3_w1_scales, + ic1, + None, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + num_m_blocks, + num_pairs, + top_k, + ) + + assert intermediate % _FP8_GROUP_SIZE == 0 + act_q = torch.empty((num_pairs, intermediate), dtype=torch.float8_e4m3fn, device=device) + act_scales = torch.empty( + (num_pairs, intermediate // _FP8_GROUP_SIZE), dtype=torch.float32, device=device + ) + if num_pairs > 0: + _swiglu_quant_fp8_kernel[(num_pairs, intermediate // _FP8_GROUP_SIZE)]( + ic1, + act_q, + act_scales, + intermediate, + ic1.stride(0), + act_q.stride(0), + act_scales.stride(0), + float(swiglu_limit) if swiglu_limit is not None else 0.0, + HAS_LIMIT=swiglu_limit is not None, + CLAMP_AFTER_SILU=swiglu_clamp_after_silu, + GROUP=_FP8_GROUP_SIZE, + FP8_MAX=_FP8_E4M3_MAX, + ) + + ic2 = torch.empty((num_pairs, hidden), dtype=torch.bfloat16, device=device) + _invoke_fp8act_fp8w_moe_kernel( + act_q, + act_scales, + w2, + w2_scales, + ic2, + token_final_scales.contiguous(), + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + num_m_blocks, + num_pairs, + 1, + ) + output = torch.empty((num_tokens, hidden), dtype=output_dtype or torch.bfloat16, device=device) + if num_tokens > 0: + block_h = 1024 + _moe_topk_sum_kernel[(num_tokens, triton.cdiv(hidden, block_h))]( + ic2, output, hidden, TOP_K=top_k, BLOCK_H=block_h + ) + return output diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index d7c3a419ccc0..9ac3fd4cc3a3 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -50,6 +50,7 @@ l0_h100: - unittest/_torch/moe/test_moe_routing.py - unittest/_torch/moe/fused_moe/test_deepgemm_fused_gather_finalize.py - unittest/_torch/moe/fused_moe/test_deepgemm_fused_expand_quant.py + - unittest/_torch/moe/fused_moe/test_triton_fp8_block_scale.py # ------------- MoE: test_moe_backend (by backend) --------------- # ------------- MoE: test_single_gpu (by backend) --------------- - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTLASS and not None" diff --git a/tests/unittest/_torch/moe/fused_moe/test_triton_fp8_block_scale.py b/tests/unittest/_torch/moe/fused_moe/test_triton_fp8_block_scale.py index 9521c34308db..83969f053168 100644 --- a/tests/unittest/_torch/moe/fused_moe/test_triton_fp8_block_scale.py +++ b/tests/unittest/_torch/moe/fused_moe/test_triton_fp8_block_scale.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ -Unit tests for ``run_triton_fp8_block_scale_moe``. +Unit tests for ``run_triton_fp8_block_scale_moe`` and ``run_triton_fp8_block_scale_moe_fp8_act``. Run as a pytest: pytest tests/unittest/_torch/moe/fused_moe/test_triton_fp8_block_scale.py -v @@ -17,7 +17,10 @@ from _torch.helpers import calc_diff, per_block_cast_to_fp8_e8m0 from tensorrt_llm._torch.moe.fused_moe.fused_moe_triton_fp8_block_scale import ( + _moe_align_small, + moe_align_block_size, run_triton_fp8_block_scale_moe, + run_triton_fp8_block_scale_moe_fp8_act, ) from tensorrt_llm._torch.moe.fused_moe.interface import ActivationType from tensorrt_llm._utils import get_sm_version @@ -201,3 +204,112 @@ def test_fp8_matches_bf16_reference(shape: MoeShape, activation_type: int) -> No diff = calc_diff(out_fp8, out_ref).item() assert diff < 1e-2, f"shape={shape.name} cosine-diff={diff:.4e} exceeds 1e-2" + + +skip_pre_hopper = pytest.mark.skipif( + not torch.cuda.is_available() or get_sm_version() < 90, + reason="FP8 Triton dot needs SM90 or newer", +) + + +def _dequant_blocks(w_fp8: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + num_experts, out_features, in_features = w_fp8.shape + expanded = scales.repeat_interleave(128, dim=1).repeat_interleave(128, dim=2) + return w_fp8.float() * expanded[:, :out_features, :in_features] + + +def _fake_quant_1x128(v: torch.Tensor) -> torch.Tensor: + """Round a 1-D FP32 vector through FP8 with one scale per 128 elements.""" + groups = v.view(-1, 128) + scale = groups.abs().amax(dim=1, keepdim=True).clamp(min=1e-10) / 448.0 + return ((groups / scale).to(torch.float8_e4m3fn).float() * scale).view(-1) + + +def _run_moe_fp8_act_reference( + x, + token_selected_experts, + token_final_scales, + w3_w1, + w3_w1_scales, + w2, + w2_scales, + swiglu_limit, + swiglu_clamp_after_silu, +) -> torch.Tensor: + """FP32 reference over dequantized weights and 1 x 128 FP8-rounded activations, + with the CUTLASS SwiGLU clamp.""" + intermediate = w3_w1.shape[1] // 2 + w31 = _dequant_blocks(w3_w1, w3_w1_scales) + w2f = _dequant_blocks(w2, w2_scales) + out = torch.zeros(x.shape, dtype=torch.float32, device=x.device) + for t in range(x.shape[0]): + x_q = _fake_quant_1x128(x[t].float()) + for k in range(token_selected_experts.shape[1]): + e = int(token_selected_experts[t, k]) + gate_up = (w31[e] @ x_q).bfloat16().float() + up, gate = gate_up[:intermediate], gate_up[intermediate:] + if swiglu_limit is not None: + up = up.clamp(-swiglu_limit, swiglu_limit) + if swiglu_clamp_after_silu: + act = F.silu(gate).clamp(max=swiglu_limit) + else: + act = F.silu(gate.clamp(max=swiglu_limit)) + else: + act = F.silu(gate) + out[t] += (w2f[e] @ _fake_quant_1x128(act * up)).bfloat16().float() * float( + token_final_scales[t, k] + ) + return out.bfloat16() + + +@skip_pre_hopper +@pytest.mark.parametrize("num_experts", [8, 72]) +@pytest.mark.parametrize("single_program", [False, True], ids=["multi_kernel", "single_program"]) +def test_moe_align_non_power_of_two_experts(num_experts: int, single_program: bool) -> None: + g = torch.Generator(device="cuda").manual_seed(3) + ids = torch.randint(0, num_experts, (5, 4), device="cuda", generator=g) + if single_program: + sorted_ids, expert_ids, num_post_pad = _moe_align_small(ids, 16, num_experts, ids.numel()) + else: + sorted_ids, expert_ids, num_post_pad = moe_align_block_size(ids, 16, num_experts) + counts = torch.bincount(ids.flatten(), minlength=num_experts) + padded = (counts + 15) // 16 * 16 + assert int(num_post_pad) == int(padded.sum()) + flat = ids.flatten() + for block in range(int(num_post_pad) // 16): + rows = sorted_ids[block * 16 : (block + 1) * 16] + valid = rows[rows < flat.numel()] + assert valid.numel() > 0 + assert bool((flat[valid] == expert_ids[block]).all()) + # Every pair lands in exactly one slot. + slots = sorted_ids[: int(num_post_pad)] + assert sorted(slots[slots < flat.numel()].tolist()) == list(range(flat.numel())) + + +@skip_pre_hopper +# 24 tokens x top-8 exceeds the single-program alignment and takes the multi-kernel one. +@pytest.mark.parametrize("num_tokens", [1, 4, 16, 24]) +@pytest.mark.parametrize( + "swiglu_limit, swiglu_clamp_after_silu", + [(None, False), (0.5, False), (0.5, True)], + ids=["no_clamp", "clamp", "clamp_after_silu"], +) +def test_fp8_act_matches_reference(num_tokens, swiglu_limit, swiglu_clamp_after_silu) -> None: + # A non-power-of-two expert count, as in decode-time routing of large MoEs. + shape = MoeShape( + "decode", num_tokens=num_tokens, hidden=512, intermediate=256, num_experts=72, top_k=8 + ) + inputs = _make_inputs(shape, seed=num_tokens) + del inputs["activation_type"] + + out = run_triton_fp8_block_scale_moe_fp8_act( + **inputs, swiglu_limit=swiglu_limit, swiglu_clamp_after_silu=swiglu_clamp_after_silu + ) + ref = _run_moe_fp8_act_reference( + **inputs, swiglu_limit=swiglu_limit, swiglu_clamp_after_silu=swiglu_clamp_after_silu + ) + + assert out.shape == ref.shape == (num_tokens, shape.hidden) + assert out.dtype == torch.bfloat16 + diff = calc_diff(out, ref).item() + assert diff < 1e-4, f"cosine-diff={diff:.4e}" From 1c7ca27d7b69966eb145d71b1abc9bc9d46bd104 Mon Sep 17 00:00:00 2001 From: Ruocheng Jia Date: Thu, 22 Oct 2026 11:52:27 -0700 Subject: [PATCH 06/14] [None][perf] Speed up FP8 block-scale MoE row expansion and activation 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 --- .../kernels/moe/cutlass/moe_kernels.cu | 29 ++++++++++++++----- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/moe/cutlass/moe_kernels.cu b/cpp/tensorrt_llm/kernels/moe/cutlass/moe_kernels.cu index 1b4959485862..909311caa19f 100644 --- a/cpp/tensorrt_llm/kernels/moe/cutlass/moe_kernels.cu +++ b/cpp/tensorrt_llm/kernels/moe/cutlass/moe_kernels.cu @@ -1862,7 +1862,17 @@ __global__ void expandInputRowsFp8BlockScaleKernel(InputActivationsType const* u int64_t const num_elems_in_col = hidden_size / ELEM_PER_THREAD; auto* const permuted_output = fp8_block_scale_out.fp8_out; - for (int64_t permuted_row = blockIdx.x; permuted_row < num_valid_tokens; permuted_row += gridDim.x) + // Each block takes a contiguous run of permuted rows: they are sorted by expert, so one search finds the + // first row's expert and the rest advance linearly instead of searching per row. + int64_t const rows_per_block = ceilDiv(num_valid_tokens, static_cast(gridDim.x)); + int64_t const row_begin = blockIdx.x * rows_per_block; + int64_t const row_end + = row_begin + rows_per_block < num_valid_tokens ? row_begin + rows_per_block : num_valid_tokens; + int64_t expert = row_begin < row_end + ? findTotalEltsLessThanTarget(expert_first_token_offset, num_experts_per_node, row_begin + 1) - 1 + : 0; + + for (int64_t permuted_row = row_begin; permuted_row < row_end; ++permuted_row) { int64_t const unpermuted_row = permuted_row_to_unpermuted_row[permuted_row]; int64_t const source_k_rank = unpermuted_row / num_tokens; @@ -1871,8 +1881,10 @@ __global__ void expandInputRowsFp8BlockScaleKernel(InputActivationsType const* u auto const* source_row_ptr = reinterpret_cast(unpermuted_input + source_row * hidden_size); auto* dest_row_ptr = permuted_output + permuted_row * hidden_size; - int64_t const expert - = findTotalEltsLessThanTarget(expert_first_token_offset, num_experts_per_node, permuted_row + 1) - 1; + while (expert + 1 < num_experts_per_node && expert_first_token_offset[expert + 1] <= permuted_row) + { + ++expert; + } int64_t const scale_row = permuted_row + fp8BlockScaleExpertPad(expert_first_token_offset[expert], expert); for (int64_t elem_index = threadIdx.x; elem_index < num_elems_in_col; elem_index += EXPAND_THREADS_PER_BLOCK) @@ -2873,8 +2885,11 @@ void doActivation(T* output, GemmOutputType const* gemm_result, float const* fp8 }(); static int32_t const sm_count = tensorrt_llm::common::getMultiProcessorCount(); - int32_t const max_blocks_per_sm - = tensorrt_llm::common::getMaxActiveBlocksPerSM(fn, ACTIVATION_THREADS_PER_BLOCK, 0); + // A row narrower than the block (e.g. a tensor-parallel slice of the intermediate size) would leave the + // remaining threads idle, so shrink the block to the row width, in whole warps. + int32_t const threads = static_cast( + std::min(ACTIVATION_THREADS_PER_BLOCK, ceilDiv(num_elems_in_col, int64_t{32}) * 32)); + int32_t const max_blocks_per_sm = tensorrt_llm::common::getMaxActiveBlocksPerSM(fn, threads, 0); // X dimension for tokens in groups of num_rows_per_cta_v // Y dimension for columns // Each CTA processes num_rows_per_cta_v * threads_x tokens @@ -2882,9 +2897,7 @@ void doActivation(T* output, GemmOutputType const* gemm_result, float const* fp8 // Divide by rows_per_cta to get the number of token blocks needed int64_t const max_token_blocks = (max_tokens + num_rows_per_cta_v - 1) / num_rows_per_cta_v; int32_t const grid_x = std::min(sm_count * max_blocks_per_sm, static_cast(max_token_blocks)); - int32_t const grid_y = static_cast( - (num_elems_in_col + ACTIVATION_THREADS_PER_BLOCK - 1) / ACTIVATION_THREADS_PER_BLOCK); - int32_t const threads = ACTIVATION_THREADS_PER_BLOCK; + int32_t const grid_y = static_cast((num_elems_in_col + threads - 1) / threads); cudaLaunchConfig_t config; config.gridDim = dim3(grid_x, grid_y, 1); From 9d9d742d528ce57542afd66b172f3695f9ce864e Mon Sep 17 00:00:00 2001 From: Ruocheng Jia Date: Thu, 22 Oct 2026 12:42:48 -0700 Subject: [PATCH 07/14] [None][perf] Avoid a device sync per KDA layer in FLA chunked prefill 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 --- tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py | 5 ++++- .../modules/kimi_kda/kimi_k3_mamba_metadata.py | 12 ++++++++++++ .../_torch/modules/kimi_kda/kimi_kda_mixer.py | 6 +++++- .../modules/kimi_kda/test_kimi_k3_mamba_metadata.py | 4 ++++ 4 files changed, 25 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py index 8d81fb9fe6b0..6768cedf784e 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py @@ -793,6 +793,7 @@ def prefill_chunk_kda( cu_seqlens: Optional[torch.Tensor], chunk_indices: Optional[torch.Tensor] = None, chunk_size: int = 64, + cu_seqlens_cpu: Optional[torch.Tensor] = None, state_pool: Optional[torch.Tensor] = None, state_indices: Optional[torch.Tensor] = None, varlen_is_aligned: Optional[bool] = None, @@ -814,7 +815,8 @@ def prefill_chunk_kda( ``chunk_indices``, ``varlen_is_aligned``, and ``single_sequence_length`` are prepared once from Kimi K3 runtime metadata. They keep the optimized path from reading ``cu_seqlens`` - back from the GPU in every KDA layer. + back from the GPU in every KDA layer; ``cu_seqlens_cpu`` (the host copy + of ``cu_seqlens``) does the same for the FLA path. """ use_indexed_state = state_pool is not None if use_indexed_state and (initial_state is not None or state_indices is None): @@ -897,6 +899,7 @@ def prefill_chunk_kda( lower_bound=lower_bound, state_v_first=True, cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, ) def decode_kda(self, **kwargs) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/modules/kimi_kda/kimi_k3_mamba_metadata.py b/tensorrt_llm/_torch/modules/kimi_kda/kimi_k3_mamba_metadata.py index 3f15c77647f0..0fd65ab9f841 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/kimi_k3_mamba_metadata.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/kimi_k3_mamba_metadata.py @@ -50,6 +50,12 @@ def __init__(self, max_batch_size: int, chunk_size: int, max_num_tokens: int) -> self.kda_chunk_indices: torch.Tensor | None = None self.kda_varlen_is_aligned: bool | None = None self.kda_single_sequence_length: int | None = None + # Host copy of the context cu_seqlens: FLA's chunked prefill derives its chunk indices from it on the + # CPU instead of reading the counts back from the GPU in every KDA layer. + self.kda_cu_seqlens_cpu: torch.Tensor | None = None + # One device cu_seqlens view per forward: FLA caches its chunk metadata by tensor identity, so every + # KDA layer must pass the same object for the cache to hit. + self.kda_cu_seqlens: torch.Tensor | None = None # KDA's CuTe verify kernel requires a 16-byte-aligned DLPack pointer. # Keep the backing allocation stable so metadata preparation remains # safe for CUDA graph capture and replay. @@ -89,8 +95,14 @@ def prepare(self, attn_metadata: AttentionMetadata) -> None: if not context_lengths: self.kda_chunk_indices = None + self.kda_cu_seqlens_cpu = None + self.kda_cu_seqlens = None return + self.kda_cu_seqlens = self.query_start_loc_long[: len(context_lengths) + 1] + + self.kda_cu_seqlens_cpu = torch.tensor([0] + context_lengths, dtype=torch.long).cumsum_(0) + chunk_counts = [ (length + KDA_PREFILL_CHUNK_SIZE - 1) // KDA_PREFILL_CHUNK_SIZE for length in context_lengths diff --git a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py index 30231196346d..0bb4cb510982 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py @@ -469,7 +469,9 @@ def _forward_impl( num_tokens = attn_metadata.num_tokens batch_size = attn_metadata.seq_lens.shape[0] state_indices = mamba_metadata.state_indices[:batch_size] - cu_seqlens = mamba_metadata.query_start_loc_long[: num_prefills + 1] + cu_seqlens = getattr(mamba_metadata, "kda_cu_seqlens", None) + if cu_seqlens is None or cu_seqlens.shape[0] != num_prefills + 1: + cu_seqlens = mamba_metadata.query_start_loc_long[: num_prefills + 1] num_decodes = batch_size - num_prefills layer_cache = attn_metadata.kv_cache_manager.mamba_layer_cache(self.layer_idx) @@ -665,6 +667,7 @@ def forward_prefill( chunk_indices = getattr(mamba_metadata, "kda_chunk_indices", None) varlen_is_aligned = getattr(mamba_metadata, "kda_varlen_is_aligned", None) single_sequence_length = getattr(mamba_metadata, "kda_single_sequence_length", None) + cu_seqlens_cpu = getattr(mamba_metadata, "kda_cu_seqlens_cpu", None) from einops import rearrange x = x2d.unsqueeze(0) # [1, T, hidden] @@ -759,6 +762,7 @@ def forward_prefill( safe_gate=lower_bound is not None, lower_bound=lower_bound, cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, chunk_indices=chunk_indices, state_pool=kernel_pool if use_indexed_state else None, state_indices=kernel_indices if use_indexed_state else None, diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py b/tests/unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py index 2553fe8f23ce..6582c07e4291 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py @@ -55,6 +55,10 @@ def test_k3_prepare_metadata_match_chunk_indices() -> None: cu_seqlens = metadata.query_start_loc_long[: len(context_lengths) + 1] legacy_chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size=64) torch.testing.assert_close(metadata.kda_chunk_indices, legacy_chunk_indices) + # The host copy FLA's chunked prefill reads instead of the device cu_seqlens. + assert metadata.kda_cu_seqlens_cpu.device.type == "cpu" + torch.testing.assert_close(metadata.kda_cu_seqlens_cpu, cu_seqlens.cpu()) + torch.testing.assert_close(metadata.kda_cu_seqlens, cu_seqlens) def test_k3_prepare_materializes_stable_aligned_generation_indices() -> None: From 85c71ebcfff7f204e6dde64e13f105b0c93dc277 Mon Sep 17 00:00:00 2001 From: Ruocheng Jia Date: Thu, 22 Oct 2026 14:31:12 -0700 Subject: [PATCH 08/14] [None][perf] Run KDA speculative verification in one launch per layer 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 --- .../_torch/modules/kimi_kda/kimi_kda_mixer.py | 140 ++++++++++++++++++ .../kimi_kda/test_kimi_kda_verify_parity.py | 8 +- 2 files changed, 147 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py index 0bb4cb510982..c1e0e74ece94 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py @@ -35,6 +35,7 @@ from ...utils import get_model_extra_attrs from ..linear import Linear from ..mamba.causal_conv1d import causal_conv1d_fn +from ..mamba.causal_conv1d_triton import causal_conv1d_update as causal_conv1d_update_triton from ..mamba.fuse_elementwise_ops import extract_transpose_prefill_slice from ..mamba.layernorm_gated import RMSNorm, rms_norm_gated_token_major from ..mamba.recurrent_state_cache import reset_recurrent_state_rows @@ -117,6 +118,21 @@ def _kda_split_conv_sections( ) +def _adjacent_sections_view(*sections: torch.Tensor) -> Optional[torch.Tensor]: + """``torch.cat(sections, -1)`` as a view when the sections already sit back to + back in one buffer (strided slices of a fused projection), else None.""" + first = sections[0] + offset = first.storage_offset() + for section in sections: + if section.stride() != first.stride() or section.storage_offset() != offset: + return None + if section.untyped_storage().data_ptr() != first.untyped_storage().data_ptr(): + return None + offset += section.shape[-1] + shape = (*first.shape[:-1], sum(section.shape[-1] for section in sections)) + return first.as_strided(shape, first.stride(), first.storage_offset()) + + def _kda_expand_fla_conv_cache(conv_state: torch.Tensor) -> torch.Tensor: """Left-pad a production ``W - 1`` cache to FLA's full ``W`` window.""" return torch.nn.functional.pad(conv_state, (1, 0)) @@ -307,6 +323,7 @@ def projection(name: str, in_features: int, out_features: int) -> nn.Module: # Output buffer for the inplace-only ``trtllm::kda_decode`` op. self._o_dense: Optional[torch.Tensor] = None self._packed_conv_weight: Optional[torch.Tensor] = None + self._verify_index_cache: dict = {} self._mtp_conv_weights: Optional[Tuple[torch.Tensor, ...]] = None self.register_to_config = False @@ -1402,6 +1419,20 @@ def forward_verify_sequential( "speculative verification requires the cache manager's " "SpeculativeState (legacy intermediate-buffer path)" ) + if self._packed_conv_weight is not None and intermediate_ssm.dtype in ( + torch.float32, + ssm_pool.dtype, + ): + return self._forward_verify_batched( + x2d, + num_steps, + intermediate_conv, + intermediate_ssm, + conv_pool, + ssm_pool, + slot_indices, + output=output, + ) d = self.proj_size num_generations = x2d.shape[0] // num_steps @@ -1477,6 +1508,115 @@ def forward_verify_sequential( # Can be removed once attention writes the core in place. return self._store_core(core, output) + def _forward_verify_batched( + self, + x2d, + num_steps, + intermediate_conv, + intermediate_ssm, + conv_pool, + ssm_pool, + slot_indices, + output: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """All verify steps in one convolution and one recurrent launch. + + Same contract as the per-step loop of :meth:`forward_verify_sequential`: + the live pools are only read, and step ``t``'s convolution window and + recurrent state land in ``intermediate_conv[:B, t]`` and + ``intermediate_ssm[:B, t]`` for ``update_mamba_states()``. + """ + from einops import rearrange + from fla.ops.kda.fused_recurrent import fused_recurrent_kda_fwd + + d = self.proj_size + num_generations = x2d.shape[0] // num_steps + num_tokens = num_generations * num_steps + device = x2d.device + x = x2d.view(num_generations, num_steps, -1) # [B, T, hidden] + + projections = self._project_verify_inputs(x, x2d.shape[0], preserve_section_strides=True) + if projections is None: + qkv = torch.cat((self.q_proj(x), self.k_proj(x), self.v_proj(x)), dim=-1) + g = self.f_b_proj(self.f_a_proj(x)) + beta = self.b_proj(x) + onorm_g = None + else: + q_proj_states, k_proj_states, v_proj_states, g, beta, onorm_g = projections + qkv = _adjacent_sections_view(q_proj_states, k_proj_states, v_proj_states) + if qkv is None: + qkv = torch.cat((q_proj_states, k_proj_states, v_proj_states), dim=-1) + if onorm_g is not None: + onorm_g = onorm_g.contiguous() + + batch_rows, state_rows, cu_seqlens = self._verify_index_tensors( + num_generations, num_steps, intermediate_ssm.shape[1], device + ) + # The convolution kernel slides the window it is given, so it runs on a + # gathered copy of the rows; per-step windows go to intermediate_conv. + conv_rows = conv_pool.index_select(0, slot_indices.long()) + conv_out = causal_conv1d_update_triton( + qkv.transpose(1, 2), # [B, 3D, T] + conv_rows, + self._packed_conv_weight, + None, + "silu", + conv_state_indices=batch_rows, + intermediate_conv_window=intermediate_conv, + intermediate_state_indices=batch_rows, + ) + q, k, v = fused_kda_post_conv( + conv_out.transpose(1, 2).reshape(num_tokens, 3 * d).t(), + num_heads=self.num_heads, + head_dim=self.head_dim, + ) + + # Step t of request b reads its state from row (b, 0), seeded with the + # pool state, and writes it to row (b, t) of intermediate_ssm. + intermediate_ssm[:num_generations, 0].copy_(ssm_pool.index_select(0, slot_indices.long())) + o, _ = fused_recurrent_kda_fwd( + q=q, + k=k, + v=v, + g=rearrange(g, "b t (h d) -> 1 (b t) h d", d=self.head_dim).contiguous(), + beta=beta.reshape(1, num_tokens, -1).contiguous(), + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=intermediate_ssm.view(-1, *intermediate_ssm.shape[2:]), + scale=self.head_k_dim**-0.5, + inplace_final_state=True, + state_v_first=True, + cu_seqlens=cu_seqlens, + ssm_state_indices=state_rows, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + lower_bound=self.gate_lower_bound, + ) + core = self._output_gate(x, o.view(num_generations, num_steps, *o.shape[2:]), onorm_g) + # Can be removed once attention writes the core in place. + return self._store_core(core, output) + + def _verify_index_tensors( + self, num_generations: int, num_steps: int, t_max: int, device: torch.device + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Batch rows, ``[B, T]`` intermediate-state rows and ``cu_seqlens`` for + :meth:`_forward_verify_batched`, built once per shape so CUDA graph + replays do not rerun them.""" + key = (num_generations, num_steps, t_max, device) + cached = self._verify_index_cache.get(key) + if cached is None: + batch_rows = torch.arange(num_generations, device=device, dtype=torch.int32) + state_rows = ( + batch_rows[:, None] * t_max + + torch.arange(num_steps, device=device, dtype=torch.int32)[None, :] + ) + cu_seqlens = torch.arange( + 0, num_generations * num_steps + 1, num_steps, device=device, dtype=torch.long + ) + cached = (batch_rows, state_rows, cu_seqlens) + self._verify_index_cache[key] = cached + return cached + def _output_gate( self, x: torch.Tensor, o: torch.Tensor, onorm_g: Optional[torch.Tensor] = None ) -> torch.Tensor: diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py index 5a06177217bb..e6487ab38d31 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py @@ -335,7 +335,8 @@ def test_kda_decode_matches_reference(num_heads): @pytest.mark.parametrize("batch", [1, 3]) @pytest.mark.parametrize("t_steps", [2, 3]) -def test_kda_verify_matches_sequential_decode(batch, t_steps): +@pytest.mark.parametrize("batched", [False, True], ids=["per_step", "batched"]) +def test_kda_verify_matches_sequential_decode(batch, t_steps, batched): if not torch.cuda.is_available(): pytest.skip("needs a GPU") torch.manual_seed(0) @@ -351,6 +352,11 @@ def test_kda_verify_matches_sequential_decode(batch, t_steps): # load_weights(); with random weights it holds heap garbage, and a # NaN/Inf bit pattern poisons both paths identically (nvbug 6599150). torch.nn.init.normal_(runtime.dt_bias, std=0.1) + if batched: + # The packed convolution weight selects the one-launch verify path. + runtime._build_mtp_conv_weights() + else: + runtime._packed_conv_weight = None slots = batch + 2 # non-trivial slot mapping cache = _LayerCache(slots, 3 * dim, w - 1, h, lin["head_dim"], lin["head_dim"], t_steps, device) slot_indices = torch.arange(2, 2 + batch, device=device, dtype=torch.long) From fd79edab7811a5d0d20a0fa759c48836f1f0b711 Mon Sep 17 00:00:00 2001 From: Ruocheng Jia Date: Thu, 22 Oct 2026 16:39:59 -0700 Subject: [PATCH 09/14] [None][perf] Fuse KDA input projections without the optimized decode 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 --- tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py index c1e0e74ece94..926bed5b30be 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py @@ -348,10 +348,11 @@ def finalize_decode_weights(self) -> None: 2. Kernel-layout constants that ``_decode_via_optimized`` used to rebuild with ~6 device kernels per layer per decode step: transposed conv weights (bf16 ``[W, D]``) and fp32 copies of - ``A_log`` / ``dt_bias`` / ``o_norm.weight``. + ``A_log`` / ``dt_bias`` / ``o_norm.weight``; only built when the + optimized decode kernel is in use. + + The fused projections serve every decode and verify path. """ - if self._dispatch.decode_kernel_path != "optimized": - return if self.q_proj.weight.device.type != "cuda": return with torch.no_grad(): @@ -381,7 +382,8 @@ def finalize_decode_weights(self) -> None: self.qkvg_proj = self._fuse_checkpoint_projections(qkvg_modules) self.qkvg_split_sizes = [module.out_features for module in qkvg_modules] self._bfa_proj_weight = self._merge_projection_weights(bfa_modules, pad_rows_to=8) - self._build_decode_kernel_constants() + if self._dispatch.decode_kernel_path == "optimized": + self._build_decode_kernel_constants() @staticmethod def _fuse_checkpoint_projections(modules: tuple[Linear, ...]) -> Linear: From 34aeb50fd774eb8ced3f726274e54f562bd8c8de Mon Sep 17 00:00:00 2001 From: Ruocheng Jia Date: Mon, 19 Oct 2026 10:31:08 -0700 Subject: [PATCH 10/14] [None][perf] Score GLM-5.3-Flash context indexer pools in FP8 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 --- .../backends/sparse/glm_kpool/backend.py | 29 +++++++- .../backends/sparse/glm_kpool/kernels.py | 69 +++++++++++++++++++ .../sparse/glm_kpool/test_kernels.py | 59 ++++++++++++++++ 3 files changed, 155 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py index 97225eb61c2a..cf2d3c5464a4 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py @@ -27,11 +27,12 @@ from __future__ import annotations +import os from dataclasses import dataclass import torch -from tensorrt_llm._utils import get_sm_version +from tensorrt_llm._utils import get_sm_version, is_sm_100f from ...interface import ( AttentionForwardArgs, @@ -46,6 +47,7 @@ gather_fp8_kv_rows, kpool_expand, kpool_score, + kpool_score_fp8, kpool_update, merge_split_attention, ) @@ -222,6 +224,9 @@ def __init__( #: construction). A float lets the kernel take it as an argument, without a #: per-call device multiply. self._native_bmm1_scale: float | None = None + #: Packed context rows score pools in FP8 with DeepGEMM MQA logits on SM100-family + #: GPUs; TRTLLM_GLM_INDEXER_BF16=1 keeps the BF16 Triton scorer for them. + self._indexer_fp8 = is_sm_100f() and os.environ.get("TRTLLM_GLM_INDEXER_BF16", "0") != "1" @classmethod def support_fused_rope(cls) -> bool: @@ -405,12 +410,32 @@ def score_pools( Incomplete/invisible pools receive the FP32 minimum before top-k. Eager callers that know a host bound on the visible lengths pass max_kv_len so only reachable pools are scored; columns past it are left unwritten and - the top-k must scan only each row's complete pools. + the top-k must scan only each row's complete pools. Packed context rows + with such a bound use the FP8 scorer on SM100-family GPUs, whose result + may be a strided view. """ state = self._cache_state(metadata) tables, gen_lens = self._rows(state, rows_per_request, request_ids=request_ids) capacity = state.block_tables.shape[1] * state.tokens_per_block kpool = self.sparse_params.index_kpool + if self._indexer_fp8 and request_ids is not None and max_kv_len is not None: + scores = kpool_score_fp8( + q, + weights, + state.index_pool, + # Packed context rows only reference the context requests' tables. + tables[: state.num_contexts], + kv_lens, + state.tokens_per_block, + head_dim=self.sparse_params.index_head_dim, + kpool=kpool, + q_scale=q_scale, + w_scale=w_scale, + request_ids=request_ids, + max_kv_len=max_kv_len, + ) + if scores is not None: + return scores return kpool_score( q, weights, diff --git a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/kernels.py b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/kernels.py index 2a6c56c21283..9f620343931c 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/kernels.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/kernels.py @@ -512,6 +512,75 @@ def kpool_score( return out +# Largest FP8 logits transient (elements) kpool_score_fp8 allocates. +_FP8_SCORE_MAX_ELEMENTS = 1 << 29 + + +def kpool_score_fp8( + q: torch.Tensor, + weights: torch.Tensor, + index_pool: torch.Tensor, + block_tables: torch.Tensor, + kv_lens: torch.Tensor, + tokens_per_block: int, + *, + head_dim: int, + kpool: int, + q_scale: float, + w_scale: float, + request_ids: torch.Tensor, + max_kv_len: int, +) -> torch.Tensor | None: + """FP8 :func:`kpool_score` for packed context rows, via DeepGEMM ``fp8_mqa_logits``. + + The pool keys of every block table are gathered into ``[tables * pools, head_dim]`` + and quantized per row; queries are quantized per (row, head), with their scales + folded into the head weights, so each score is ``sum_h w * relu(q . key)`` as in + :func:`kpool_score`. ``max_kv_len`` bounds every row's ``kv_lens``. Row ``r`` + scores only the first ``kv_lens[r] // kpool`` pools of table ``request_ids[r]``; + later columns are not meaningful, so readers must stop there. The result may be + a strided view with an aligned row stride. Returns ``None`` when the logits + transient would exceed ``_FP8_SCORE_MAX_ELEMENTS``. + """ + from tensorrt_llm.deep_gemm import fp8_mqa_logits + + n = q.shape[0] + pools = max(0, int(max_kv_len)) // kpool + tables = block_tables.shape[0] + if n == 0 or pools == 0 or n * tables * pools > _FP8_SCORE_MAX_ELEMENTS: + return None + positions = torch.arange(pools, device=q.device) * kpool + page = positions // tokens_per_block + offset = positions - page * tokens_per_block + slots = block_tables[:, page].long() # [tables, pools] + keys = index_pool[slots, offset, 2 * head_dim : 3 * head_dim].reshape(-1, head_dim) + half = head_dim // 2 + keys_fp8, key_scale = torch.ops.trtllm.fused_cat_fp8(keys[:, :half], keys[:, half:], False) + q_rows = q.reshape(-1, head_dim) + q_fp8, q_row_scale = torch.ops.trtllm.fused_cat_fp8(q_rows[:, :half], q_rows[:, half:], False) + head_weights = weights.float() * q_row_scale.view(n, -1) * (w_scale * q_scale) + start = request_ids.to(torch.int32) * pools + end = start + (kv_lens // kpool).to(torch.int32) + logits = fp8_mqa_logits( + q_fp8.view(q.shape), + (keys_fp8, key_scale.view(-1)), + head_weights, + start, + end, + clean_logits=False, + ) + if tables == 1: + # DeepGEMM pads the logits rows to an aligned stride, as the pool top-k needs. + return logits[:, :pools] + # Keep the copy's rows aligned and wider than the selection, like the logits arena. + out = torch.empty( + n, max(1024, triton.cdiv(pools, 64) * 64), dtype=torch.float32, device=q.device + ) + rows = torch.arange(n, device=q.device) + out[:, :pools] = logits[:, : tables * pools].view(n, tables, pools)[rows, request_ids.long()] + return out + + @triton.jit def _kpool_expand_kernel( SEL, diff --git a/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py b/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py index b13dab5487ce..1575a2d02ee5 100644 --- a/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py +++ b/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py @@ -22,6 +22,7 @@ compact_sparse_rows, kpool_expand, kpool_score, + kpool_score_fp8, kpool_update, merge_split_attention, ) @@ -321,6 +322,64 @@ def test_pool_scores_broadcast_table_single_context_request(): ) +@pytest.mark.skipif( + get_sm_version() not in (100, 103), reason="DeepGEMM MQA logits require SM100/103" +) +@pytest.mark.parametrize("single_request", [True, False], ids=["one_table", "two_tables"]) +def test_fp8_pool_scores_track_bf16_selection(single_request): + """FP8 context scoring keeps the BF16 scorer's top pools on packed context rows.""" + gen = torch.Generator(device="cuda").manual_seed(10) + heads, tpb, slots = 32, 32, 400 + lens = [6000] if single_request else [6000, 2500] + pages = [(length + tpb - 1) // tpb for length in lens] + pool = torch.randn(slots, 2, tpb, 3 * HD, generator=gen, device="cuda").to(torch.bfloat16)[:, 1] + perm = torch.randperm(slots, generator=gen, device="cuda") + tables = torch.zeros(len(lens), max(pages), dtype=torch.int64, device="cuda") + start = 0 + for i, count in enumerate(pages): + tables[i, :count] = perm[start : start + count] + start += count + request_ids = torch.cat( + [torch.full((length,), i, dtype=torch.int32) for i, length in enumerate(lens)] + ).cuda() + kv_lens = torch.cat([torch.arange(1, length + 1, dtype=torch.int32) for length in lens]).cuda() + n = kv_lens.shape[0] + q = torch.randn(n, heads, HD, generator=gen, device="cuda").to(torch.bfloat16) + w = torch.randn(n, heads, generator=gen, device="cuda").to(torch.bfloat16) + kwargs = dict( + head_dim=HD, + kpool=KPOOL, + q_scale=HD**-0.5, + w_scale=heads**-0.5, + request_ids=request_ids, + max_kv_len=max(lens), + ) + want = kpool_score( + q, + w, + pool, + tables, + kv_lens, + tpb, + num_pools_max=max(pages) * tpb // KPOOL, + precision="bf16", + rows_per_program=16, + **kwargs, + ) + got = kpool_score_fp8(q, w, pool, tables, kv_lens, tpb, **kwargs) + assert got is not None and got.stride(1) == 1 and got.stride(0) % 4 == 0 + overlaps = [] + for row in torch.randint(0, n, (64,), generator=gen, device="cuda").tolist(): + visible = int(kv_lens[row]) // KPOOL + if visible < 64: + continue + k = visible // 2 + a = set(want[row, :visible].topk(k).indices.tolist()) + b = set(got[row, :visible].topk(k).indices.tolist()) + overlaps.append(len(a & b) / k) + assert overlaps and sum(overlaps) / len(overlaps) > 0.97 and min(overlaps) > 0.9 + + def test_pool_expansion_graph_replay_tracks_request_tables_and_tails(): # Requests occupy disjoint slots, with gaps between slots from coalesced buffers. tables = torch.tensor([[3, 1], [5, 2]], dtype=torch.int64, device="cuda") From ed36729db40960ab1db5806d515cc3425b28b3cb Mon Sep 17 00:00:00 2001 From: Ruocheng Jia Date: Tue, 20 Oct 2026 07:38:35 -0700 Subject: [PATCH 11/14] [None][fix] Bound GLM-5.3-Flash context pool-scoring memory 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 --- .../backends/sparse/glm_kpool/backend.py | 89 ++++++++++++++++++- .../_torch/models/modeling_glm5_next.py | 72 ++++++++++----- tensorrt_llm/_torch/pyexecutor/_util.py | 11 ++- .../sparse/glm_kpool/test_glm_kpool.py | 82 +++++++++++++++++ .../sparse/glm_kpool/test_kernels.py | 45 ++++++++++ 5 files changed, 269 insertions(+), 30 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py index cf2d3c5464a4..d585c1e01cef 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py @@ -28,6 +28,7 @@ from __future__ import annotations import os +from collections.abc import Sequence from dataclasses import dataclass import torch @@ -143,6 +144,27 @@ class GlmKpoolSparseAttention(TrtllmAttention): # At the published 2112-entry selection this uses 132 MiB of BF16 rows. _FP8_QUERY_CHUNK_SIZE = 64 + # FP32 elements one packed-context scoring call may hold: logits, their row + # copy and gathered pool keys grow with the visible pools, which a fresh-prompt + # profiling forward does not reach. plan_context_scores tiles calls to this + # bound and runtime_workspace_fixed_bytes reserves it in KV-cache sizing. + _CONTEXT_SCORE_ELEMENTS = 1 << 26 + + @classmethod + def runtime_workspace_fixed_bytes( + cls, + model_config, + mapping, + *, + max_batch_size: int, + max_num_tokens: int, + max_seq_len: int, + enable_cuda_graph: bool, + ) -> int: + if max_num_tokens <= 0 or max_seq_len <= 0: + return 0 + return min(cls._CONTEXT_SCORE_ELEMENTS, max_num_tokens * max_seq_len) * 4 + def __init__( self, layer_idx: int, @@ -389,6 +411,62 @@ def update_pool_keys( request_ids=request_ids, ) + def plan_context_scores( + self, metadata, spans: Sequence[tuple[int, int]] + ) -> list[tuple[int, int, int, int, int]]: + """Split packed context scoring into calls within ``_CONTEXT_SCORE_ELEMENTS``. + + spans lists ``(rows, visible bound)`` of each context request in packed + order. Each returned ``(row_start, row_end, table_start, table_end, + max_kv_len)`` scores rows ``row_start:row_end`` against context tables + ``table_start:table_end``. Consecutive requests share a call while it + fits; a request that does not fit alone is split by rows. Costs count + only the transients that grow with the visible pools. + """ + state = self._cache_state(metadata) + kpool = self.sparse_params.index_kpool + budget = self._CONTEXT_SCORE_ELEMENTS + # The BF16 scorer writes every pool column of the configured capacity. + capacity_pools = -(-state.block_tables.shape[1] * state.tokens_per_block // kpool) + # Gathered BF16 pool keys plus their FP8 copy, in FP32 elements. + key_elements = self.sparse_params.index_head_dim * 3 // 4 + + def cost(rows: int, tables: int, pools: int) -> int: + if not self._indexer_fp8 or pools == 0: + return rows * capacity_pools + logits = -(-tables * pools // 256) * 256 + copy = max(1024, -(-pools // 64) * 64) if tables > 1 else 0 + return rows * (logits + copy) + tables * pools * key_elements + + calls: list[tuple[int, int, int, int, int]] = [] + group: list[int] | None = None # [row_start, rows, table_start, tables, max visible] + row_start = 0 + for table, (rows, visible) in enumerate(spans): + if group is not None: + merged = max(group[4], visible) + if cost(group[1] + rows, group[3] + 1, merged // kpool) <= budget: + group[1] += rows + group[3] += 1 + group[4] = merged + row_start += rows + continue + calls.append( + (group[0], group[0] + group[1], group[2], group[2] + group[3], group[4]) + ) + group = None + if cost(rows, 1, visible // kpool) <= budget: + group = [row_start, rows, table, 1, visible] + else: + fixed = cost(0, 1, visible // kpool) + tile = max(1, (budget - fixed) // max(1, cost(1, 1, visible // kpool) - fixed)) + for start in range(row_start, row_start + rows, tile): + end = min(start + tile, row_start + rows) + calls.append((start, end, table, table + 1, visible)) + row_start += rows + if group is not None: + calls.append((group[0], group[0] + group[1], group[2], group[2] + group[3], group[4])) + return calls + def score_pools( self, q: torch.Tensor, @@ -401,6 +479,7 @@ def score_pools( rows_per_request: int = 1, request_ids: torch.Tensor | None = None, max_kv_len: int | None = None, + table_range: tuple[int, int] | None = None, ) -> torch.Tensor: """Score complete pools into FP32 [rows, pool_capacity]. @@ -412,26 +491,28 @@ def score_pools( only reachable pools are scored; columns past it are left unwritten and the top-k must scan only each row's complete pools. Packed context rows with such a bound use the FP8 scorer on SM100-family GPUs, whose result - may be a strided view. + may be a strided view. table_range, from :meth:`plan_context_scores`, + names the context tables those rows' request_ids fall in. """ state = self._cache_state(metadata) tables, gen_lens = self._rows(state, rows_per_request, request_ids=request_ids) capacity = state.block_tables.shape[1] * state.tokens_per_block kpool = self.sparse_params.index_kpool if self._indexer_fp8 and request_ids is not None and max_kv_len is not None: + # Packed context rows only reference the context requests' tables. + first, last = (0, state.num_contexts) if table_range is None else table_range scores = kpool_score_fp8( q, weights, state.index_pool, - # Packed context rows only reference the context requests' tables. - tables[: state.num_contexts], + tables[first:last], kv_lens, state.tokens_per_block, head_dim=self.sparse_params.index_head_dim, kpool=kpool, q_scale=q_scale, w_scale=w_scale, - request_ids=request_ids, + request_ids=request_ids - first if first else request_ids, max_kv_len=max_kv_len, ) if scores is not None: diff --git a/tensorrt_llm/_torch/models/modeling_glm5_next.py b/tensorrt_llm/_torch/models/modeling_glm5_next.py index dcc02d2098ab..dc2aca4e58a8 100644 --- a/tensorrt_llm/_torch/models/modeling_glm5_next.py +++ b/tensorrt_llm/_torch/models/modeling_glm5_next.py @@ -1272,13 +1272,18 @@ def _select_and_attend_paged( reduce: bool = True, request_ids: torch.Tensor | None = None, max_visible: int | None = None, + num_candidates: torch.Tensor | None = None, + context_spans: Sequence[tuple[int, int]] | None = None, ) -> torch.Tensor: """Score pools, select members and tail, then attend to paged latent rows. visible gives each query's prefix length. Packed context uses request_ids; verification uses rows_per_request. Single-token decode reads live lengths from metadata. All device work uses fixed buffer shapes for graph replay. - max_visible is an optional host bound on visible that limits scoring work. + max_visible is an optional host bound on visible that limits scoring work; + num_candidates, if given, is visible // index_kpool as int32. Packed context + passes context_spans, each request's (rows, visible bound), so scoring + runs in calls of bounded memory. """ indexer = self.indexer rows = q_resid.shape[0] @@ -1287,26 +1292,42 @@ def _select_and_attend_paged( # each request contributes several rows (context or verification). plain_decode = request_ids is None and rows_per_request == 1 kv_lens = None if plain_decode else visible - scores = self.attn_backend.score_pools( - q_index, - index_weights, - metadata, - q_scale=indexer.softmax_scale, - w_scale=indexer.head_mix_scale, - kv_lens=kv_lens, - rows_per_request=rows_per_request, - request_ids=request_ids, - max_kv_len=max_visible, - ) - num_cand = (visible // indexer.index_kpool).to(torch.int32) - selected = torch.empty(rows, indexer.select_k, dtype=torch.int32, device=scores.device) - indexer.pool_top_k( - scores, - selected, - is_prefill=False, - sequence_lengths=num_cand, - scan_lengths=num_cand, + num_cand = ( + (visible // indexer.index_kpool).to(torch.int32) + if num_candidates is None + else num_candidates ) + selected = torch.empty(rows, indexer.select_k, dtype=torch.int32, device=q_index.device) + # Each call scores and selects its own rows; scores are dropped before the next. + if context_spans is None: + calls = [(slice(None), max_visible, None)] + else: + calls = [ + (slice(start, end), max_kv_len, (first, last)) + for start, end, first, last, max_kv_len in self.attn_backend.plan_context_scores( + metadata, context_spans + ) + ] + for part, max_kv_len, table_range in calls: + scores = self.attn_backend.score_pools( + q_index[part], + index_weights[part], + metadata, + q_scale=indexer.softmax_scale, + w_scale=indexer.head_mix_scale, + kv_lens=None if kv_lens is None else kv_lens[part], + rows_per_request=rows_per_request, + request_ids=None if request_ids is None else request_ids[part], + max_kv_len=max_kv_len, + table_range=table_range, + ) + indexer.pool_top_k( + scores, + selected[part], + is_prefill=False, + sequence_lengths=num_cand[part], + scan_lengths=num_cand[part], + ) topk_rows = self.attn_backend.expand_selection( selected, metadata, @@ -1362,10 +1383,14 @@ def forward_prefill( request_ids=rows.final_request_ids, ) query = self.q_b_proj(q_resid).view(-1, self.num_heads, self.qk_head_dim) - max_visible = max( - int(cached_lens[i]) + int(cu_seqlens[i + 1]) - int(cu_seqlens[i]) + spans = [ + ( + int(cu_seqlens[i + 1]) - int(cu_seqlens[i]), + int(cached_lens[i]) + int(cu_seqlens[i + 1]) - int(cu_seqlens[i]), + ) for i in range(len(cu_seqlens) - 1) - ) + ] + max_visible = max(visible for _, visible in spans) return self._select_and_attend_paged( q_resid, query, @@ -1376,6 +1401,7 @@ def forward_prefill( reduce=reduce, request_ids=rows.request_ids, max_visible=max_visible, + context_spans=spans, ) def forward_decode( diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 72b0e12823ab..b1c167e375f1 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -547,9 +547,14 @@ def get_attention_workspace_fixed_bytes(model_config, mapping, **capacities) -> int: """Capacity-sized reserve declared by the model's selected backend.""" from ..attention.backends.utils import get_attention_backend - return get_attention_backend( - model_config.attn_backend).runtime_workspace_fixed_bytes( - model_config, mapping, **capacities) + backend = get_attention_backend(model_config.attn_backend) + if config_utils.is_glm5_next(model_config.pretrained_config): + # glm5_next layers select the glm_kpool sparse backend themselves, so + # model_config carries no sparse_attention_config naming it. + from ..attention.backends.sparse.glm_kpool import \ + GlmKpoolSparseAttention as backend + return backend.runtime_workspace_fixed_bytes(model_config, mapping, + **capacities) def get_overlap_headroom_penalty_workspace_bytes( diff --git a/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py b/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py index c1bbb6d5143b..cb7466aa8159 100644 --- a/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py +++ b/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py @@ -501,3 +501,85 @@ def test_latent_pool_rows_rejects_layouts_without_a_uniform_row_view(): ) with pytest.raises(ValueError, match="not multiples of dim"): latent_pool_rows(odd_stride) + + +def _planner(fp8: bool, pages: int = 4096, tokens_per_block: int = 32): + backend = object.__new__(GlmKpoolSparseAttention) + backend.sparse_params = GlmKpoolSparseParams() + backend._indexer_fp8 = fp8 + state = SimpleNamespace( + block_tables=torch.zeros(8, pages, dtype=torch.long), tokens_per_block=tokens_per_block + ) + backend._cache_state = Mock(return_value=state) + return backend + + +def _check_plan(calls, spans, budget, fp8): + """Calls cover every row in order, name each row's table, and fit the FP8 logits budget.""" + request_starts = [0] + for rows, _ in spans: + request_starts.append(request_starts[-1] + rows) + assert calls[0][0] == 0 and calls[-1][1] == request_starts[-1] + for (_, end, _, _, _), (start, _, _, _, _) in zip(calls, calls[1:]): + assert end == start + kpool = GlmKpoolSparseParams().index_kpool + for start, end, first, last, max_kv_len in calls: + assert start < end and first < last + assert request_starts[first] <= start and end <= request_starts[last] + assert max_kv_len >= max(visible for _, visible in spans[first:last]) + if fp8: + assert (end - start) * (last - first) * (max_kv_len // kpool) <= budget + + +@pytest.mark.parametrize( + "spans", + [ + [(13324, 13324)], # one fresh long prompt + [(512, 16384 + 512)] * 4, # cached prefixes + [(16384, 131072)], # last chunk of a 128K prompt + [(3, 3)] * 64 + [(4000, 4000)], # tiny prompts beside a long one + ], + ids=["fresh", "prefixes", "long-chunk", "tiny-and-long"], +) +@pytest.mark.parametrize("fp8", [False, True], ids=["bf16", "fp8"]) +def test_context_score_plan_stays_within_the_reserved_bound(spans, fp8): + backend = _planner(fp8) + budget = GlmKpoolSparseAttention._CONTEXT_SCORE_ELEMENTS + calls = backend.plan_context_scores(object(), spans) + _check_plan(calls, spans, budget, fp8) + if not fp8: + capacity_pools = 4096 * 32 // GlmKpoolSparseParams().index_kpool + assert all((end - start) * capacity_pools <= budget for start, end, *_ in calls) + + +def test_context_score_plan_keeps_a_fresh_prompt_in_one_call(): + calls = _planner(True).plan_context_scores(object(), [(13324, 13324)]) + assert calls == [(0, 13324, 0, 1, 13324)] + + +def test_context_score_plan_groups_short_requests_and_splits_a_long_one(): + backend = _planner(True) + grouped = backend.plan_context_scores(object(), [(100, 100)] * 3) + assert grouped == [(0, 300, 0, 3, 100)] + split = backend.plan_context_scores(object(), [(16384, 131072)]) + assert len(split) > 1 and {call[2:] for call in split} == {(0, 1, 131072)} + + +def test_kv_cache_sizing_reserves_the_context_score_bound(): + from tensorrt_llm._torch.pyexecutor._util import get_attention_workspace_fixed_bytes + + capacities = dict( + max_batch_size=64, max_num_tokens=16384, max_seq_len=131072, enable_cuda_graph=True + ) + + def reserve(model_type): + config = SimpleNamespace( + attn_backend="TRTLLM", + sparse_attention_config=None, + pretrained_config=SimpleNamespace(model_type=model_type), + ) + return get_attention_workspace_fixed_bytes(config, None, **capacities) + + bound = GlmKpoolSparseAttention._CONTEXT_SCORE_ELEMENTS * 4 + assert reserve("glm5_next") == reserve("glm5_next_text") == bound + assert reserve("llama") == 0 diff --git a/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py b/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py index 1575a2d02ee5..f3e0923dd846 100644 --- a/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py +++ b/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py @@ -380,6 +380,51 @@ def test_fp8_pool_scores_track_bf16_selection(single_request): assert overlaps and sum(overlaps) / len(overlaps) > 0.97 and min(overlaps) > 0.9 +@pytest.mark.skipif( + get_sm_version() not in (100, 103), reason="DeepGEMM MQA logits require SM100/103" +) +def test_fp8_pool_scores_of_a_table_range_match_the_packed_call(): + """Scoring one request against its own table equals its rows of a multi-table call.""" + gen = torch.Generator(device="cuda").manual_seed(11) + heads, tpb, slots = 32, 32, 400 + lens = [3000, 2500, 1200] + pages = [(length + tpb - 1) // tpb for length in lens] + pool = torch.randn(slots, 2, tpb, 3 * HD, generator=gen, device="cuda").to(torch.bfloat16)[:, 1] + perm = torch.randperm(slots, generator=gen, device="cuda") + tables = torch.zeros(len(lens), max(pages), dtype=torch.int64, device="cuda") + start = 0 + for i, count in enumerate(pages): + tables[i, :count] = perm[start : start + count] + start += count + request_ids = torch.cat( + [torch.full((length,), i, dtype=torch.int32) for i, length in enumerate(lens)] + ).cuda() + kv_lens = torch.cat([torch.arange(1, length + 1, dtype=torch.int32) for length in lens]).cuda() + q = torch.randn(kv_lens.shape[0], heads, HD, generator=gen, device="cuda").to(torch.bfloat16) + w = torch.randn(kv_lens.shape[0], heads, generator=gen, device="cuda").to(torch.bfloat16) + kwargs = dict(head_dim=HD, kpool=KPOOL, q_scale=HD**-0.5, w_scale=heads**-0.5) + packed = kpool_score_fp8( + q, w, pool, tables, kv_lens, tpb, request_ids=request_ids, max_kv_len=max(lens), **kwargs + ) + first = lens[0] + rows = slice(first, first + lens[1]) + alone = kpool_score_fp8( + q[rows], + w[rows], + pool, + tables[1:2], + kv_lens[rows], + tpb, + request_ids=request_ids[rows] - 1, + max_kv_len=lens[1], + **kwargs, + ) + assert packed is not None and alone is not None + for row in range(0, lens[1], 97): + visible = (row + 1) // KPOOL + torch.testing.assert_close(alone[row, :visible], packed[first + row, :visible]) + + def test_pool_expansion_graph_replay_tracks_request_tables_and_tails(): # Requests occupy disjoint slots, with gaps between slots from coalesced buffers. tables = torch.tensor([[3, 1], [5, 2]], dtype=torch.int64, device="cuda") From 2470636c80b72a8a1f47c8beec2fb313166cc552 Mon Sep 17 00:00:00 2001 From: Ruocheng Jia Date: Thu, 22 Oct 2026 00:16:01 -0700 Subject: [PATCH 12/14] [None][fix] Count every GLM-5.3-Flash context-scoring buffer in its bound 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 --- .../backends/sparse/glm_kpool/backend.py | 55 ++++++++++++++--- .../_torch/models/modeling_glm5_next.py | 4 +- .../sparse/glm_kpool/test_glm_kpool.py | 59 ++++++++++++++++--- .../sparse/glm_kpool/test_kernels.py | 50 ++++++++++++++++ 4 files changed, 150 insertions(+), 18 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py index d585c1e01cef..937b4aec77e6 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py @@ -33,6 +33,7 @@ import torch +from tensorrt_llm._torch.pyexecutor.config_utils import unwrap_glm5_next_text_config from tensorrt_llm._utils import get_sm_version, is_sm_100f from ...interface import ( @@ -150,6 +151,27 @@ class GlmKpoolSparseAttention(TrtllmAttention): # bound and runtime_workspace_fixed_bytes reserves it in KV-cache sizing. _CONTEXT_SCORE_ELEMENTS = 1 << 26 + @staticmethod + def _context_score_cost( + rows: int, tables: int, pools: int, *, fp8: bool, head_dim: int, capacity_pools: int + ) -> int: + """FP32 elements a scoring call of ``rows`` rows over ``tables`` tables holds at once. + + Counts the buffers that grow with the visible ``pools``: the BF16 scorer's + capacity-wide output, or every buffer :func:`kpool_score_fp8` keeps alive + together -- the padded logits, a multi-table call's gather temporary and + aligned row copy, and per gathered key its BF16 row, FP8 row, scale and + int64 slot index. + """ + if not fp8 or pools == 0: + return rows * capacity_pools + # fp8_mqa_logits pads each row to a multiple of 256 columns, plus 256. + logits = -(-tables * pools // 256) * 256 + 256 + copy = max(1024, -(-pools // 64) * 64) + pools if tables > 1 else 0 + keys = tables * pools * (head_dim // 2 + head_dim // 4 + 3) + # int64 pool positions, pages and offsets. + return rows * (logits + copy) + keys + 6 * pools + @classmethod def runtime_workspace_fixed_bytes( cls, @@ -161,9 +183,25 @@ def runtime_workspace_fixed_bytes( max_seq_len: int, enable_cuda_graph: bool, ) -> int: + """Bytes of the largest scoring call plan_context_scores can emit for these capacities.""" if max_num_tokens <= 0 or max_seq_len <= 0: return 0 - return min(cls._CONTEXT_SCORE_ELEMENTS, max_num_tokens * max_seq_len) * 4 + config = unwrap_glm5_next_text_config(model_config.pretrained_config) + defaults = GlmKpoolSparseParams() + kpool = getattr(config, "index_kpool", defaults.index_kpool) + head_dim = getattr(config, "index_head_dim", defaults.index_head_dim) + pools = -(-max_seq_len // kpool) + tables = max(1, min(max_batch_size, max_num_tokens)) + + def cost(rows: int, tables: int, fp8: bool) -> int: + return cls._context_score_cost( + rows, tables, pools, fp8=fp8, head_dim=head_dim, capacity_pools=pools + ) + + largest = max(cost(max_num_tokens, tables, True), cost(max_num_tokens, 1, False)) + # Calls fit the budget, except that a single row is never split. + bound = max(cls._CONTEXT_SCORE_ELEMENTS, cost(1, 1, True), cost(1, 1, False)) + return min(largest, bound) * 4 def __init__( self, @@ -428,15 +466,16 @@ def plan_context_scores( budget = self._CONTEXT_SCORE_ELEMENTS # The BF16 scorer writes every pool column of the configured capacity. capacity_pools = -(-state.block_tables.shape[1] * state.tokens_per_block // kpool) - # Gathered BF16 pool keys plus their FP8 copy, in FP32 elements. - key_elements = self.sparse_params.index_head_dim * 3 // 4 def cost(rows: int, tables: int, pools: int) -> int: - if not self._indexer_fp8 or pools == 0: - return rows * capacity_pools - logits = -(-tables * pools // 256) * 256 - copy = max(1024, -(-pools // 64) * 64) if tables > 1 else 0 - return rows * (logits + copy) + tables * pools * key_elements + return self._context_score_cost( + rows, + tables, + pools, + fp8=self._indexer_fp8, + head_dim=self.sparse_params.index_head_dim, + capacity_pools=capacity_pools, + ) calls: list[tuple[int, int, int, int, int]] = [] group: list[int] | None = None # [row_start, rows, table_start, tables, max visible] diff --git a/tensorrt_llm/_torch/models/modeling_glm5_next.py b/tensorrt_llm/_torch/models/modeling_glm5_next.py index dc2aca4e58a8..87ec4fc7bc14 100644 --- a/tensorrt_llm/_torch/models/modeling_glm5_next.py +++ b/tensorrt_llm/_torch/models/modeling_glm5_next.py @@ -1298,7 +1298,7 @@ def _select_and_attend_paged( else num_candidates ) selected = torch.empty(rows, indexer.select_k, dtype=torch.int32, device=q_index.device) - # Each call scores and selects its own rows; scores are dropped before the next. + # Each call scores and selects its own rows. if context_spans is None: calls = [(slice(None), max_visible, None)] else: @@ -1328,6 +1328,8 @@ def _select_and_attend_paged( sequence_lengths=num_cand[part], scan_lengths=num_cand[part], ) + # Release this call's scores before the next call allocates its own. + del scores topk_rows = self.attn_backend.expand_selection( selected, metadata, diff --git a/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py b/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py index cb7466aa8159..c98dd52f4994 100644 --- a/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py +++ b/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py @@ -514,21 +514,37 @@ def _planner(fp8: bool, pages: int = 4096, tokens_per_block: int = 32): return backend -def _check_plan(calls, spans, budget, fp8): - """Calls cover every row in order, name each row's table, and fit the FP8 logits budget.""" +def _call_costs(backend, calls): + """Each planned call's FP32-element cost under the backend's cost model.""" + state = backend._cache_state(None) + kpool = backend.sparse_params.index_kpool + capacity_pools = state.block_tables.shape[1] * state.tokens_per_block // kpool + return [ + GlmKpoolSparseAttention._context_score_cost( + end - start, + last - first, + max_kv_len // kpool, + fp8=backend._indexer_fp8, + head_dim=backend.sparse_params.index_head_dim, + capacity_pools=capacity_pools, + ) + for start, end, first, last, max_kv_len in calls + ] + + +def _check_plan(backend, calls, spans, budget): + """Calls cover every row in order, name each row's table, and fit the budget.""" request_starts = [0] for rows, _ in spans: request_starts.append(request_starts[-1] + rows) assert calls[0][0] == 0 and calls[-1][1] == request_starts[-1] for (_, end, _, _, _), (start, _, _, _, _) in zip(calls, calls[1:]): assert end == start - kpool = GlmKpoolSparseParams().index_kpool for start, end, first, last, max_kv_len in calls: assert start < end and first < last assert request_starts[first] <= start and end <= request_starts[last] assert max_kv_len >= max(visible for _, visible in spans[first:last]) - if fp8: - assert (end - start) * (last - first) * (max_kv_len // kpool) <= budget + assert max(_call_costs(backend, calls)) <= budget @pytest.mark.parametrize( @@ -546,10 +562,7 @@ def test_context_score_plan_stays_within_the_reserved_bound(spans, fp8): backend = _planner(fp8) budget = GlmKpoolSparseAttention._CONTEXT_SCORE_ELEMENTS calls = backend.plan_context_scores(object(), spans) - _check_plan(calls, spans, budget, fp8) - if not fp8: - capacity_pools = 4096 * 32 // GlmKpoolSparseParams().index_kpool - assert all((end - start) * capacity_pools <= budget for start, end, *_ in calls) + _check_plan(backend, calls, spans, budget) def test_context_score_plan_keeps_a_fresh_prompt_in_one_call(): @@ -565,6 +578,34 @@ def test_context_score_plan_groups_short_requests_and_splits_a_long_one(): assert len(split) > 1 and {call[2:] for call in split} == {(0, 1, 131072)} +@pytest.mark.parametrize( + "capacities", + [ + dict(max_batch_size=64, max_num_tokens=64, max_seq_len=4096), + dict(max_batch_size=8, max_num_tokens=2048, max_seq_len=32768), + dict(max_batch_size=64, max_num_tokens=16384, max_seq_len=131072), + ], + ids=["small-token-budget", "mid", "large"], +) +@pytest.mark.parametrize("fp8", [False, True], ids=["bf16", "fp8"]) +def test_kv_cache_reserve_covers_every_planned_call(capacities, fp8): + tokens, seq_len = capacities["max_num_tokens"], capacities["max_seq_len"] + requests = min(capacities["max_batch_size"], tokens) + config = SimpleNamespace(pretrained_config=SimpleNamespace(model_type="glm5_next_text")) + reserve = GlmKpoolSparseAttention.runtime_workspace_fixed_bytes( + config, None, enable_cuda_graph=True, **capacities + ) + backend = _planner(fp8, pages=seq_len // 32) + batches = [ + [(1, seq_len)] * requests, # one continuation row per request + [(tokens, seq_len)], # one request takes the whole token budget + [(tokens // requests, seq_len)] * requests, # the budget split across requests + ] + for spans in batches: + calls = backend.plan_context_scores(object(), spans) + assert max(_call_costs(backend, calls)) * 4 <= reserve + + def test_kv_cache_sizing_reserves_the_context_score_bound(): from tensorrt_llm._torch.pyexecutor._util import get_attention_workspace_fixed_bytes diff --git a/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py b/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py index f3e0923dd846..5951232e6592 100644 --- a/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py +++ b/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py @@ -380,6 +380,56 @@ def test_fp8_pool_scores_track_bf16_selection(single_request): assert overlaps and sum(overlaps) / len(overlaps) > 0.97 and min(overlaps) > 0.9 +@pytest.mark.skipif( + get_sm_version() not in (100, 103), reason="DeepGEMM MQA logits require SM100/103" +) +@pytest.mark.parametrize("lens", [[16384], [16000, 12000, 8000]], ids=["one-table", "three-tables"]) +def test_fp8_pool_score_peak_memory_fits_the_planned_cost(lens): + """kpool_score_fp8's peak allocation beyond its row-sized buffers fits the planner's cost.""" + from tensorrt_llm._torch.attention.backends.sparse.glm_kpool import GlmKpoolSparseAttention + + gen = torch.Generator(device="cuda").manual_seed(13) + heads, tpb, rows_per_table = 32, 32, 256 + pages = [(length + tpb - 1) // tpb for length in lens] + pool = torch.randn(sum(pages) + 8, 2, tpb, 3 * HD, generator=gen, device="cuda") + pool = pool.to(torch.bfloat16)[:, 1] + tables = torch.zeros(len(lens), max(pages), dtype=torch.int64, device="cuda") + start = 0 + for i, count in enumerate(pages): + tables[i, :count] = torch.arange(start, start + count, device="cuda") + start += count + # The last rows_per_table queries of each prompt, so most pools are visible. + request_ids = torch.cat( + [torch.full((rows_per_table,), i, dtype=torch.int32) for i in range(len(lens))] + ).cuda() + kv_lens = torch.cat( + [ + torch.arange(length - rows_per_table + 1, length + 1, dtype=torch.int32) + for length in lens + ] + ).cuda() + n = kv_lens.shape[0] + q = torch.randn(n, heads, HD, generator=gen, device="cuda").to(torch.bfloat16) + w = torch.randn(n, heads, generator=gen, device="cuda").to(torch.bfloat16) + kwargs = dict(head_dim=HD, kpool=KPOOL, q_scale=HD**-0.5, w_scale=heads**-0.5) + args = (q, w, pool, tables, kv_lens, tpb) + kpool_score_fp8(*args, request_ids=request_ids, max_kv_len=max(lens), **kwargs) # JIT warmup + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + base = torch.cuda.memory_allocated() + scores = kpool_score_fp8(*args, request_ids=request_ids, max_kv_len=max(lens), **kwargs) + torch.cuda.synchronize() + peak = torch.cuda.max_memory_allocated() - base + del scores + cost = GlmKpoolSparseAttention._context_score_cost( + n, len(lens), max(lens) // KPOOL, fp8=True, head_dim=HD, capacity_pools=0 + ) + # Quantized queries (with the quantizer's scratch), head weights and per-row + # indices grow only with rows; startup profiling measures them. + row_bytes = n * heads * (2 * HD + 16) + 64 * n + assert peak <= cost * 4 + row_bytes + (64 << 10) + + @pytest.mark.skipif( get_sm_version() not in (100, 103), reason="DeepGEMM MQA logits require SM100/103" ) From 1cb9865ea6afce5ea91397045d4e5d4ea65c2865 Mon Sep 17 00:00:00 2001 From: Ruocheng Jia Date: Thu, 22 Oct 2026 20:18:16 -0700 Subject: [PATCH 13/14] [None][perf] Score GLM-5.3-Flash context indexer pools in FP8 on SM90 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 --- .../attention/backends/sparse/glm_kpool/backend.py | 10 ++++++---- .../_torch/attention/sparse/glm_kpool/test_kernels.py | 6 +++--- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py index 937b4aec77e6..e8f963137556 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py @@ -284,9 +284,11 @@ def __init__( #: construction). A float lets the kernel take it as an argument, without a #: per-call device multiply. self._native_bmm1_scale: float | None = None - #: Packed context rows score pools in FP8 with DeepGEMM MQA logits on SM100-family - #: GPUs; TRTLLM_GLM_INDEXER_BF16=1 keeps the BF16 Triton scorer for them. - self._indexer_fp8 = is_sm_100f() and os.environ.get("TRTLLM_GLM_INDEXER_BF16", "0") != "1" + #: Packed context rows score pools in FP8 with DeepGEMM MQA logits on SM90 and + #: SM100-family GPUs; TRTLLM_GLM_INDEXER_BF16=1 keeps the BF16 Triton scorer for them. + self._indexer_fp8 = (get_sm_version() == 90 or is_sm_100f()) and os.environ.get( + "TRTLLM_GLM_INDEXER_BF16", "0" + ) != "1" @classmethod def support_fused_rope(cls) -> bool: @@ -529,7 +531,7 @@ def score_pools( callers that know a host bound on the visible lengths pass max_kv_len so only reachable pools are scored; columns past it are left unwritten and the top-k must scan only each row's complete pools. Packed context rows - with such a bound use the FP8 scorer on SM100-family GPUs, whose result + with such a bound use the FP8 scorer on SM90 and SM100-family GPUs, whose result may be a strided view. table_range, from :meth:`plan_context_scores`, names the context tables those rows' request_ids fall in. """ diff --git a/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py b/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py index 5951232e6592..9e9cd442a292 100644 --- a/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py +++ b/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py @@ -323,7 +323,7 @@ def test_pool_scores_broadcast_table_single_context_request(): @pytest.mark.skipif( - get_sm_version() not in (100, 103), reason="DeepGEMM MQA logits require SM100/103" + get_sm_version() not in (90, 100, 103), reason="DeepGEMM MQA logits require SM90 or SM100/103" ) @pytest.mark.parametrize("single_request", [True, False], ids=["one_table", "two_tables"]) def test_fp8_pool_scores_track_bf16_selection(single_request): @@ -381,7 +381,7 @@ def test_fp8_pool_scores_track_bf16_selection(single_request): @pytest.mark.skipif( - get_sm_version() not in (100, 103), reason="DeepGEMM MQA logits require SM100/103" + get_sm_version() not in (90, 100, 103), reason="DeepGEMM MQA logits require SM90 or SM100/103" ) @pytest.mark.parametrize("lens", [[16384], [16000, 12000, 8000]], ids=["one-table", "three-tables"]) def test_fp8_pool_score_peak_memory_fits_the_planned_cost(lens): @@ -431,7 +431,7 @@ def test_fp8_pool_score_peak_memory_fits_the_planned_cost(lens): @pytest.mark.skipif( - get_sm_version() not in (100, 103), reason="DeepGEMM MQA logits require SM100/103" + get_sm_version() not in (90, 100, 103), reason="DeepGEMM MQA logits require SM90 or SM100/103" ) def test_fp8_pool_scores_of_a_table_range_match_the_packed_call(): """Scoring one request against its own table equals its rows of a multi-table call.""" From 5aa96d61d79d9f03a895bd94fa81f9994d4fda30 Mon Sep 17 00:00:00 2001 From: Ruocheng Jia Date: Fri, 23 Oct 2026 01:38:07 -0700 Subject: [PATCH 14/14] [None][perf] Support breakable prefill CUDA graphs for GLM-5.3-Flash 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 --- .../_torch/models/modeling_glm5_next.py | 82 +++++++++++++++++-- .../modeling/test_glm5_next_contracts.py | 42 ++++++++++ 2 files changed, 116 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_glm5_next.py b/tensorrt_llm/_torch/models/modeling_glm5_next.py index 87ec4fc7bc14..b9639016daa4 100644 --- a/tensorrt_llm/_torch/models/modeling_glm5_next.py +++ b/tensorrt_llm/_torch/models/modeling_glm5_next.py @@ -23,6 +23,7 @@ import copy import re +import weakref from collections.abc import Iterable, Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Callable @@ -58,8 +59,9 @@ from ..modules.linear import Linear, TensorParallelMode from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm +from ..pyexecutor.breakable_cuda_graph import eager_on_graph, is_in_breakable_cuda_graph from ..pyexecutor.config_utils import unwrap_glm5_next_text_config -from ..utils import AuxStreamType +from ..utils import AuxStreamType, get_model_extra_attrs from .checkpoints.hf.glm5_next_weight_mapper import ( Glm5NextHfWeightMapper, glm5_next_is_quantized, @@ -85,6 +87,20 @@ SPARSE_MLP = "sparse" +def _register_graph_layer( + model_config: ModelConfig, registry: str, layer_idx: int, module: nn.Module +) -> str: + """Register ``module`` under a unique key for breakable-CUDA-graph eager ops.""" + layers = model_config.extra_attrs.setdefault(registry, {}) + key = str(layer_idx) + suffix = 0 + while key in layers: + key = f"{layer_idx}_{suffix}" + suffix += 1 + layers[key] = weakref.ref(module) + return key + + @dataclass(frozen=True) class Glm5NextSchedule: """The two literal per-layer dispatch lists, validated against each other.""" @@ -879,6 +895,7 @@ def __init__( mapping: Mapping | None = None, allreduce_strategy: AllReduceStrategy = AllReduceStrategy.ONESHOT, aux_stream: torch.cuda.Stream | None = None, + model_config: ModelConfig | None = None, ) -> None: del dtype # the shared mixer is bf16 (fp32 gate parameters and pools) # The checkpoint's ``linear_attn_config`` does not name the gate rank: @@ -898,6 +915,12 @@ def __init__( allreduce_strategy=allreduce_strategy, aux_stream=aux_stream, ) + # Join the shared mixer's breakable-CUDA-graph registry. model_config is + # not passed to the mixer itself, whose projections keep this module's + # quantization setup. + if model_config is not None: + self.layer_idx_str = _register_graph_layer(model_config, "kda_layers", layer_idx, self) + self.register_to_config = True # Preserve checkpoint FP32 gate parameters. MetaInitMode rejects detach # on meta tensors, so create empty FP32 parameters on that branch. for name in ("A_log", "dt_bias"): @@ -1164,6 +1187,11 @@ def __init__( self.indexer = Glm5NextIndexer( config, layer_idx, mapping=mapping, model_config=model_config ) + self._graph_key = ( + _register_graph_layer(model_config, "glm5_sparse_layers", layer_idx, self) + if model_config is not None + else None + ) # The production sparse-MLA backend, selected through the standard # dispatch (`get_attention_backend(attn_backend, sparse_params)` via @@ -1454,12 +1482,25 @@ def forward( Reuse the model's per-forward schedules and row cache across layers. Reduce the complete output once, including context/verify mixtures. + Inside a breakable prefill CUDA graph the phase split runs eagerly + between graph segments over the live rows of the padded batch. """ - ctx = ( - runtime_ctx - if runtime_ctx is not None - else build_glm5_next_runtime_context(attn_metadata) - ) + if is_in_breakable_cuda_graph() and self._graph_key is not None: + out = hidden_states.new_empty((hidden_states.shape[0], self.hidden_size)) + maybe_bcg_glm5_sparse_attention_inplace(hidden_states, self._graph_key, out) + else: + ctx = ( + runtime_ctx + if runtime_ctx is not None + else build_glm5_next_runtime_context(attn_metadata) + ) + out = self._forward_phases(hidden_states, ctx) + return out if self.tp_all_reduce is None else self.tp_all_reduce(out) + + def _forward_phases( + self, hidden_states: torch.Tensor, ctx: Glm5NextRuntimeContext + ) -> torch.Tensor: + """Run each batch phase and return the unreduced ``[tokens, hidden]`` output.""" parts = [] if ctx.num_contexts: parts.append( @@ -1488,8 +1529,7 @@ def forward( reduce=False, ) parts.append(out) - out = parts[0] if len(parts) == 1 else torch.cat(parts, dim=0) - return out if self.tp_all_reduce is None else self.tp_all_reduce(out) + return parts[0] if len(parts) == 1 else torch.cat(parts, dim=0) def forward_verify( self, @@ -1545,6 +1585,27 @@ def forward_verify( ) +def glm5_sparse_attention_inplace( + hidden_states: torch.Tensor, layer_key: str, output: torch.Tensor +) -> None: + """Run one sparse-attention layer's phases on the live rows into ``output``. + + Rows past the live token count (CUDA-graph padding) are zeroed. + """ + extra_attrs = get_model_extra_attrs() + metadata = extra_attrs["attention_metadata"]() + layer = extra_attrs["glm5_sparse_layers"][layer_key]() + num_tokens = metadata.num_tokens + out = layer._forward_phases( + hidden_states[:num_tokens], build_glm5_next_runtime_context(metadata) + ) + output[:num_tokens].copy_(out) + output[num_tokens:].zero_() + + +maybe_bcg_glm5_sparse_attention_inplace = eager_on_graph(glm5_sparse_attention_inplace) + + class Glm5NextGate(DeepseekV3Gate): """DeepSeek noaux_tc routing with FP32 logits and correction bias. @@ -1799,6 +1860,7 @@ def __init__( mapping=mapping, allreduce_strategy=model_config.allreduce_strategy, aux_stream=attention_aux_stream, + model_config=model_config, ) else: self.self_attn = Glm5NextSparseAttention( @@ -1944,6 +2006,10 @@ def forward( raise ValueError("specify exactly one of input_ids or inputs_embeds") if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) + if attn_metadata.padded_num_tokens is not None and not is_in_breakable_cuda_graph(): + # Eager passes over a CUDA-graph-padded batch run only the live rows; + # the speculative wrapper trims the output to the same rows. + inputs_embeds = inputs_embeds[: attn_metadata.num_tokens] if runtime_ctx is None: runtime_ctx = build_glm5_next_runtime_context(attn_metadata) # Only the decoder stack: under one-model MTP the speculative base diff --git a/tests/unittest/_torch/modeling/test_glm5_next_contracts.py b/tests/unittest/_torch/modeling/test_glm5_next_contracts.py index 5763a76c2ed8..3bbdf6a1cdf6 100644 --- a/tests/unittest/_torch/modeling/test_glm5_next_contracts.py +++ b/tests/unittest/_torch/modeling/test_glm5_next_contracts.py @@ -1093,6 +1093,7 @@ def attend(x, *args, reduce=True, **kwargs): create_autospec(implementation, side_effect=attend), ) attention.forward = MethodType(Glm5NextSparseAttention.forward, attention) + attention._forward_phases = MethodType(Glm5NextSparseAttention._forward_phases, attention) connection = SimpleNamespace( pre_mapping=lambda streams: (None, None, streams.mean(dim=1)), post_mapping=lambda output, residual, post, comb: residual + output.unsqueeze(1), @@ -1181,6 +1182,7 @@ def test_packed_attention_splits_phases_and_reduces_once(generation_phase, reduc forward_verify=Mock(return_value=gen), tp_all_reduce=reduction, ) + attention._forward_phases = MethodType(Glm5NextSparseAttention._forward_phases, attention) metadata = SimpleNamespace() context = Glm5NextRuntimeContext( num_contexts=1, @@ -1227,3 +1229,43 @@ def test_indexer_topk_without_cute_dsl(monkeypatch: pytest.MonkeyPatch) -> None: expected = torch.arange(max(0, count - indexer.select_k), count, dtype=torch.int32) assert torch.equal(valid, expected) assert int((selected[row] == -1).sum()) == indexer.select_k - min(count, indexer.select_k) + + +@pytest.mark.cpu_only +def test_breakable_graph_registry_keys_are_unique(): + from tensorrt_llm._torch.models import modeling_glm5_next as glm + + model_config = SimpleNamespace(extra_attrs={}) + first, second = torch.nn.Linear(1, 1), torch.nn.Linear(1, 1) + keys = [ + glm._register_graph_layer(model_config, "glm5_sparse_layers", 3, module) + for module in (first, second) + ] + layers = model_config.extra_attrs["glm5_sparse_layers"] + assert keys == ["3", "3_0"] + assert layers["3"]() is first and layers["3_0"]() is second + + +@pytest.mark.cpu_only +def test_sparse_attention_graph_op_fills_live_rows_and_zeros_padding(monkeypatch): + from tensorrt_llm._torch.models import modeling_glm5_next as glm + from tensorrt_llm._torch.utils import model_extra_attrs + + class _Layer: + def _forward_phases(self, hidden_states, ctx): + assert ctx == "ctx" and hidden_states.shape[0] == 3 + return hidden_states * 2 + + layer = _Layer() + metadata = SimpleNamespace(num_tokens=3) + monkeypatch.setattr(glm, "build_glm5_next_runtime_context", lambda m: "ctx") + hidden = torch.arange(20, dtype=torch.float32).view(5, 4) + output = torch.full((5, 4), float("nan")) + attrs = { + "attention_metadata": lambda: metadata, + "glm5_sparse_layers": {"7": lambda: layer}, + } + with model_extra_attrs(attrs): + glm.glm5_sparse_attention_inplace(hidden, "7", output) + torch.testing.assert_close(output[:3], hidden[:3] * 2) + assert torch.equal(output[3:], torch.zeros(2, 4))