From aadb953bc631707f84567ba6ae90bafda1e29a25 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 15:02:44 -0700 Subject: [PATCH 01/53] Add the glm_moe_dsa modeling_v2 family with its sm107 tp16 target layout Route GlmMoeDsaForCausalLM to models/glm_moe_dsa: GLM-5.2-NVFP4 on sm107 with 16-way tensor parallelism, no attention DP and one-model MTP-Eagle (any draft length) or no speculation. Add Arch.SM_107 for entries certified on Rubin. The target's modeling and weights follow. Signed-off-by: Dongfeng Yu --- .../modeling_v2/_router_index.py | 1 + .../_experimental/modeling_v2/catalog/_op.py | 1 + .../models/glm_moe_dsa/__init__.py | 3 + .../glm_5_2_nvfp4__sm_107__tp16/__init__.py | 3 + .../glm_5_2_nvfp4__sm_107__tp16/modeling.py | 20 +++++ .../ops/__init__.py | 3 + .../ops/attention/__init__.py | 2 + .../ops/comm/__init__.py | 2 + .../ops/dsa/__init__.py | 2 + .../ops/embedding/__init__.py | 2 + .../ops/gemm/__init__.py | 2 + .../ops/moe/__init__.py | 2 + .../ops/norm/__init__.py | 2 + .../ops/spec/__init__.py | 2 + .../glm_5_2_nvfp4__sm_107__tp16/weights.py | 3 + .../modeling_v2/models/glm_moe_dsa/routing.py | 83 +++++++++++++++++++ 16 files changed, 133 insertions(+) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/embedding/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/norm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index b0b6e78ef8d8..5233f575bd73 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -54,6 +54,7 @@ # Qwen3.8-Flash-Next publishes ``Qwen4ExpForConditionalGeneration``, and the # config loader's composite branch writes the same name back. "Qwen4ExpForConditionalGeneration": "models.qwen4_exp.routing", + "GlmMoeDsaForCausalLM": "models.glm_moe_dsa.routing", } diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/_op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/_op.py index 3dc75761b4a3..4baee1ba64b6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/_op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/_op.py @@ -88,6 +88,7 @@ class Arch(str, enum.Enum): """ SM_103 = "sm_103" + SM_107 = "sm_107" @dataclass(frozen=True) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/__init__.py new file mode 100644 index 000000000000..43c6570013d1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GlmMoeDsaForCausalLM targets (GLM-5.2) and the routing between them.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/__init__.py new file mode 100644 index 000000000000..9bcf8ba8c02c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GLM-5.2-NVFP4 on 16 Rubin GPUs (sm107): tensor-parallel attention and experts over 16 ranks.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py new file mode 100644 index 000000000000..bb2cf4453939 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GLM-5.2-NVFP4 on 16 Rubin GPUs (sm107), tp16: model construction and forward.""" + +from __future__ import annotations + + +class PrefillTarget: + """Forwards that carry context requests.""" + + +class DecodeTarget: + """Generation-only forwards.""" + + +class ModelingV2Glm52Nvfp4Sm107Tp16: + """GLM-5.2-NVFP4, sm107, tp16.""" + + def __init__(self, model_config): + raise NotImplementedError("the glm_moe_dsa tp16 target is not assembled yet") diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/__init__.py new file mode 100644 index 000000000000..f8bbe8b7c789 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Entries only this target calls, one module per op under a category package.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/embedding/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/embedding/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/norm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/norm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py new file mode 100644 index 000000000000..82ab171a1966 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GLM-5.2-NVFP4 checkpoint tensors mapped onto the tp16 target's parameters.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py new file mode 100644 index 000000000000..bec0b9684538 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Where a GlmMoeDsaForCausalLM config lands. Read this file and you know. + +``route`` is one forward-reading decision: the facts that fix which code runs, +compared against the deployments the targets below were written for. When none +fits it raises ``NoTarget`` stating those facts. In ``auto`` the engine then +uses the built-in implementation (``modeling_deepseekv3``); in ``require`` it +raises that error. +""" + +from __future__ import annotations + +from typing import Any, Optional, Tuple + +from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig + +from ..._router_index import NoTarget + +# (layers, hidden, routed experts, q_lora_rank). Shape-sniffing also matches +# fine-tunes of the same shape; the accuracy gate pins identity. +_SHAPE = (78, 6144, 256, 2048) + + +def route( + pretrained_config: Any, + mapping: Any, + compute_capability: Tuple[int, int], + spec_config: Optional[Any], +) -> type: + """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits. + + Each target is imported where it is returned, so declining costs no import + of modeling code. + """ + shape = ( + pretrained_config.num_hidden_layers, + pretrained_config.hidden_size, + pretrained_config.n_routed_experts, + pretrained_config.q_lora_rank, + ) + # tp16: attention heads and the expert intermediate dimension sharded over + # 16 ranks, no attention DP, no expert parallelism. A weight-layout + # question, which is why it selects a target rather than a runtime branch. + tp16 = ( + mapping.world_size == 16 + and mapping.tp_size == 16 + and mapping.moe_ep_size == 1 + and mapping.moe_tp_size == 16 + and not mapping.enable_attention_dp + and mapping.pp_size == 1 + and mapping.cp_size == 1 + ) + # The target drives the checkpoint's one MTP module through the one-model + # MTP-Eagle loop as a linear chain, at any draft length, or runs without + # speculative decoding. + mtp_eagle_or_none = spec_config is None or ( + isinstance(spec_config, MTPDecodingConfig) + and spec_config.spec_dec_mode.is_mtp_eagle_one_model() + and not getattr(spec_config, "use_dynamic_tree", False) + and getattr(spec_config, "eagle_choices", None) is None + and spec_config.max_total_draft_tokens == spec_config.max_draft_len + ) + # GLM-5.2-NVFP4 on 16 Rubin GPUs (sm107), tp16, MTP-Eagle or no speculation. + if compute_capability == (10, 7) and shape == _SHAPE and tp16 and mtp_eagle_or_none: + # One line on purpose: the claims suite reads the routed targets off + # `from ..modeling import ` lines. + # fmt: off + # isort: off + from .glm_5_2_nvfp4__sm_107__tp16.modeling import ModelingV2Glm52Nvfp4Sm107Tp16 + # isort: on + # fmt: on + return ModelingV2Glm52Nvfp4Sm107Tp16 + raise NoTarget( + f"compute capability {compute_capability}, config shape {shape}, " + f"world_size={mapping.world_size} tp={mapping.tp_size} moe_ep={mapping.moe_ep_size} " + f"moe_tp={mapping.moe_tp_size} attention_dp={mapping.enable_attention_dp} " + f"pp={mapping.pp_size} cp={mapping.cp_size}, speculative decoding " + f"{spec_config.spec_dec_mode.name if spec_config is not None else 'off'}; " + "the one glm_moe_dsa target is GLM-5.2-NVFP4 on sm107, tp16 (16 ranks, " + "tensor-parallel attention and experts, no attention DP), with one-model " + "MTP-Eagle as a linear chain or no speculative decoding" + ) From 0cf150f9d5a1dab8ae39655839dcacb78fdd03d7 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 15:16:23 -0700 Subject: [PATCH 02/53] Pad the token-sparse MLA Q tile to 8 rows The FMHA autotuner sets tileSizeQ of the token-sparse MLA generation kernels to the Q heads per KV head. With fewer than 8 heads per rank (64 MLA heads over 16 ranks: 4) the SwapsMmaAb kernel then gets a Q tile that SmemTile rejects when it is JIT-compiled. Pad the tile to 8 rows on every SM, as the dense MLA selection does; the kernel masks the rows past the heads. Signed-off-by: Dongfeng Yu --- .../trtllmGenKernels/fmha/fmhaKernels.h | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index 79cc1eb5dd7e..80fe9fc45c99 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -675,7 +675,10 @@ class TllmGenFmhaKernel if (mSM != kSM_107) { FmhaAutoTuner autoTuner(options, optionsFromArgs, multiProcessorCount); - return autoTuner.selectKernel(); + int32_t ctaDim = 512; + std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); + padTokenSparseMlaQTile(options); + return {options, optionsFromArgs, ctaDim}; } if (isGmemReduction(options.mMultiCtasKvMode) && options.mTileScheduler == TileScheduler::Static) @@ -692,10 +695,23 @@ class TllmGenFmhaKernel int32_t ctaDim = 512; FmhaAutoTuner autoTuner(options, optionsFromArgs, multiProcessorCount); std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); + padTokenSparseMlaQTile(options); limitCgaSmemReductionCtasKv(options); return {options, optionsFromArgs, ctaDim}; } + // The token-sparse MLA selection sets tileSizeQ to the Q heads per KV head, while the SwapsMmaAb kernels need a + // multiple of 8 Q-tile rows (SmemTile). With fewer than 8 heads per rank (for example 64 MLA heads over 16 ranks) + // pad the tile to 8, as the dense MLA selection does; the kernel masks the rows past the heads. + static void padTokenSparseMlaQTile(FmhaOptions& options) + { + if (options.mIsMlaGen && isTokenSparse(options.mSparseType) + && options.mFmhaKernelType == FmhaKernelType::SwapsMmaAbForGeneration && options.mTileSizeQ < 8) + { + options.mTileSizeQ = 8; + } + } + void limitCgaSmemReductionCtasKv(FmhaOptions& options) const { if (!isCgaSmemReduction(options.mMultiCtasKvMode)) From 48bab5fea1eef287e4cf5487acf707aee65e20cb Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 15:17:38 -0700 Subject: [PATCH 03/53] Keep the dense MLP at the full TP when the post-MLP all-reduce is fused When the TP group spans nodes, _compute_mlp_tp_size capped the dense MLP at gpus_per_node, while the fused post-MLP all-reduce sums the MLP partials over the whole TP group: each partial was counted tp_size / mlp_tp_size times (2x at TP8 and 4x at TP16 on 4-GPU nodes). Cap the MLP TP only when the MLP reduces within its own group (eager fusion disabled); with the fusion on, the MLP keeps the full TP. Signed-off-by: Dongfeng Yu --- tensorrt_llm/_torch/models/modeling_deepseekv3.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index 6a54201f9065..b2bdc7cf78e7 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -1476,8 +1476,9 @@ def _get_decoder_layer_quant_config( def _compute_mlp_tp_size(self, intermediate_size: int, block_size: int) -> int: """ - For DeepSeek‑R1, MLP TP size is limited by intermediate_size // block_size - and must also be multiples of gpus_per_node to avoid expensive inter‑node allreduce. + For DeepSeek‑R1, MLP TP size is limited by intermediate_size // block_size. + Without the fused post-MLP all-reduce it is also capped at gpus_per_node to + avoid an expensive inter‑node allreduce. Args: intermediate_size (int): MLP intermediate size. @@ -1499,12 +1500,17 @@ def _compute_mlp_tp_size(self, intermediate_size: int, self.mapping.tp_size, ) - if tp > self.mapping.gpus_per_node: + if tp > self.mapping.gpus_per_node and not self.enable_fusion: + # Without the fused post-MLP all-reduce the MLP reduces within + # its own group, so a node-local MLP TP stays correct. mlp_tp_size = math.gcd( tp, self.mapping.gpus_per_node, ) # Avoid costly inter-node TP else: + # The fused post-MLP all-reduce sums over every rank of the TP + # group, so the MLP keeps the full TP: a smaller MLP TP would + # count each partial sum tp_size / mlp_tp_size times. mlp_tp_size = tp return mlp_tp_size From 6bedb704a2bffba674f9eb53c6903527aea14fef Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 15:23:12 -0700 Subject: [PATCH 04/53] Assemble the GLM-5.2 tp16 target on the built-in modules The target builds the built-in GLM-5.2 modules and loads the checkpoint with the built-in loader, so the MLA KV cache and DSA indexer K cache it reads have the layout the context server wrote. Forwards with context rows run that decoder stack in PrefillTarget; generation-only forwards run DecodeTarget, which starts from the same stack. Signed-off-by: Dongfeng Yu --- .../glm_5_2_nvfp4__sm_107__tp16/modeling.py | 194 +++++++++++++++++- .../glm_5_2_nvfp4__sm_107__tp16/weights.py | 47 ++++- 2 files changed, 234 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py index bb2cf4453939..e7e86a427b58 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py @@ -1,20 +1,202 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GLM-5.2-NVFP4 on 16 Rubin GPUs (sm107), tp16: model construction and forward.""" +"""GLM-5.2-NVFP4 on 16 Rubin GPUs (sm107), tp16: model construction and forward. + +This target serves the generation side of a disaggregated deployment. Its +context server runs the built-in GLM-5.2 implementation (modeling_deepseekv3) +and hands over the MLA KV cache and the DSA indexer K cache it wrote, so this +target builds the same modules and loads the same weights: the caches it reads +are laid out exactly as they were written. Forwards that carry context rows +(engine warm-up, an aggregated deployment) run that same decoder stack in +``PrefillTarget``. Generation-only forwards -- every serving step of the +generation server -- run ``DecodeTarget``. +""" from __future__ import annotations +from typing import Any, Literal, Optional + +import torch +from transformers import PretrainedConfig + +from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3ForCausalLM, DeepseekV3Model +from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM +from tensorrt_llm._torch.speculative import SpecMetadata + +from . import weights as _weights + +# The deployment this target is written for; route() selects it on the same +# facts, and the constructor states them again for a direct construction. +_COMPUTE_CAPABILITY = (10, 7) +_TP = 16 + class PrefillTarget: - """Forwards that carry context requests.""" + """Forwards that carry context rows: the decoder stack the context server runs.""" + + def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: + self.core = core + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, + **kwargs, + ) -> torch.Tensor: + return DeepseekV3Model.forward( + self.core, + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + **kwargs, + ) class DecodeTarget: - """Generation-only forwards.""" + """Generation-only forwards: the serving step of the generation server. + + Routed to only when the batch has no context rows, which the core decides + per step; nothing in this class reads the phase back. + """ + + def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: + self.core = core + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, + **kwargs, + ) -> torch.Tensor: + return DeepseekV3Model.forward( + self.core, + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + **kwargs, + ) + +class Glm52Nvfp4Sm107Tp16Core(DeepseekV3Model): + """The decoder stack, with the per-step choice between the two targets.""" -class ModelingV2Glm52Nvfp4Sm107Tp16: + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + super().__init__(model_config) + self._prefill: Optional[PrefillTarget] = None + self._decode: Optional[DecodeTarget] = None + + def post_load_weights(self) -> None: + """Build the targets once the weights are final. Reached more than once + (the shell and the loader's module walk), so it only builds once.""" + if self._decode is not None: + return + self._prefill = PrefillTarget(self) + self._decode = DecodeTarget(self) + + def _select_target(self, attn_metadata: AttentionMetadata) -> "PrefillTarget | DecodeTarget": + if self._decode is None: + self.post_load_weights() + return self._prefill if attn_metadata.num_contexts > 0 else self._decode + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, + **kwargs, + ) -> torch.Tensor: + return self._select_target(attn_metadata).forward( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + **kwargs, + ) + + +class ModelingV2Glm52Nvfp4Sm107Tp16( + SpecDecOneEngineForCausalLM[Glm52Nvfp4Sm107Tp16Core, PretrainedConfig] +): """GLM-5.2-NVFP4, sm107, tp16.""" - def __init__(self, model_config): - raise NotImplementedError("the glm_moe_dsa tp16 target is not assembled yet") + @classmethod + def get_preferred_kv_cache_manager_version( + cls, pretrained_config: Any = None + ) -> Literal["V2"]: + """The KV cache manager the context server's built-in implementation prefers.""" + return "V2" + + @classmethod + def get_preferred_transceiver_runtime( + cls, pretrained_config: Any = None + ) -> Optional[Literal["CPP", "PYTHON"]]: + """The transceiver runtime for the masked DSA indexer K-cache transfer.""" + return "PYTHON" + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + mapping = model_config.mapping + assert torch.cuda.get_device_capability() == _COMPUTE_CAPABILITY, ( + f"this target runs on sm107, not {torch.cuda.get_device_capability()}" + ) + assert ( + mapping.world_size == _TP + and mapping.tp_size == _TP + and mapping.moe_ep_size == 1 + and not mapping.enable_attention_dp + and mapping.pp_size == 1 + ), f"this target is tp16 without attention DP or expert parallelism, not {mapping}" + model_config = _weights.normalize_model_config(model_config) + super().__init__(model=Glm52Nvfp4Sm107Tp16Core(model_config), model_config=model_config) + if ( + model_config.spec_config is not None + and model_config.spec_config.spec_dec_mode.is_mtp_one_model() + ): + self.num_hidden_layers = self.config.num_hidden_layers + self.model.layers.extend(self.draft_model.mtp_layers) + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: torch.IntTensor = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + return_context_logits: bool = False, + spec_metadata: Optional[SpecMetadata] = None, + resource_manager=None, + **kwargs, + ) -> torch.Tensor: + return super().forward( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + return_context_logits=return_context_logits, + spec_metadata=spec_metadata, + resource_manager=resource_manager, + **kwargs, + ) + + def load_weights(self, weights, *args, **kwargs): + _weights.load_weights(self, weights) + + def setup_aliases(self) -> None: + DeepseekV3ForCausalLM.setup_aliases(self) + + def post_load_weights(self): + super().post_load_weights() + self.model.post_load_weights() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py index 82ab171a1966..1e6491eff23f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py @@ -1,3 +1,48 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GLM-5.2-NVFP4 checkpoint tensors mapped onto the tp16 target's parameters.""" +"""GLM-5.2-NVFP4 checkpoint tensors mapped onto the tp16 target's parameters. + +The target's modules are the built-in GLM-5.2 modules (see ``modeling.py``), +so the checkpoint loads through the built-in loader into exactly the shapes +and layouts the context server uses. +""" + +from __future__ import annotations + +import copy + +from transformers import PretrainedConfig + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3WeightLoader + + +def normalize_model_config(model_config: ModelConfig[PretrainedConfig]) -> ModelConfig[PretrainedConfig]: + """The two config rewrites the built-in implementation applies before building modules. + + - Legacy checkpoints name the fused q_a / kv_a projection ``fused_a`` in their + per-module quant config; the modules call it ``kv_a_proj_with_mqa``. + - GLM-5.2 publishes ``model_type = glm_moe_dsa``; its attention is DeepSeek-V3.2's + (MLA with the DSA indexer), which the modules select by ``deepseek_v32``. + """ + if model_config.quant_config_dict is not None: + model_config = copy.deepcopy(model_config) + quant_config_dict = {} + for key, val in model_config.quant_config_dict.items(): + key_split = key.split(".") + if key_split[-1] == "fused_a": + key = ".".join(key_split[:-1] + ["kv_a_proj_with_mqa"]) + quant_config_dict[key] = val + model_config._frozen = False + model_config.quant_config_dict = quant_config_dict + model_config._frozen = True + if model_config.pretrained_config.model_type == "glm_moe_dsa": + model_config._frozen = False + model_config.pretrained_config.model_type = "deepseek_v32" + model_config._frozen = True + return model_config + + +def load_weights(model, weights) -> None: + """Load the checkpoint into the target's modules with the built-in loader.""" + DeepseekV3WeightLoader(model).load_weights(weights) From 41a8b939ebe7bca3174f0f6d1df7cb44606197b4 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 15:49:14 -0700 Subject: [PATCH 05/53] Cover the glm_moe_dsa tp16 route in the routing tests The route matches GLM-5.2-NVFP4 on sm107 at tp16 without attention DP or expert parallelism, with one-model MTP-Eagle at any draft length or no speculation, and declines another shape, attention DP, expert parallelism, tp8, vanilla MTP, NGram and sm103, naming the criterion that missed. Signed-off-by: Dongfeng Yu --- .../modeling_v2/test_modeling_v2_routing.py | 71 +++++++++++++++++++ 1 file changed, 71 insertions(+) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py index 81256a1c2c97..42d1d187f016 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -338,3 +338,74 @@ def test_qwen4_exp_deepstack_vision_tower_does_not_match(): def test_qwen4_exp_on_two_ranks_does_not_match(): config = _model_config(_qwen4_exp_config(), world_size=2, tp_size=2) assert modeling_v2_resolve(config, _llm_args("auto")) is None + + +# --- glm_moe_dsa: tp16 on sm107, MTP-Eagle over a linear chain at any length ---- + +_GLM52_TARGET = ( + "models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.modeling.ModelingV2Glm52Nvfp4Sm107Tp16" +) +_TP16 = dict(world_size=16, tp_size=16, moe_ep_size=1, moe_tp_size=16) + + +@pytest.fixture +def _on_sm107(monkeypatch): + """Route as if this were a Rubin GPU.""" + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (10, 7)) + + +def _glm52_config(**overrides): + """The shape GLM-5.2-NVFP4's own config.json declares.""" + fields = dict( + architectures=["GlmMoeDsaForCausalLM"], + model_type="glm_moe_dsa", + num_hidden_layers=78, + hidden_size=6144, + n_routed_experts=256, + q_lora_rank=2048, + ) + fields.update(overrides) + return PretrainedConfig(**fields) + + +@pytest.mark.parametrize("mode", ["auto", "require"]) +@pytest.mark.parametrize( + "spec_config", + [None, MTPDecodingConfig(max_draft_len=7), MTPDecodingConfig(max_draft_len=3)], + ids=["no-spec", "mtp7", "mtp3"], +) +def test_glm52_tp16_matches(_on_sm107, spec_config, mode): + """Any draft length: the draft loop replays the one MTP module per step.""" + config = _model_config(_glm52_config(), spec_config=spec_config, **_TP16) + assert _qualified(modeling_v2_resolve(config, _llm_args(mode))) == _GLM52_TARGET + + +@pytest.mark.parametrize( + "config_kwargs, mapping_kwargs, spec_config, missed", + [ + (dict(num_hidden_layers=92), _TP16, None, r"shape \(92, 6144, 256, 2048\)"), + (dict(), dict(_TP16, enable_attention_dp=True), None, "attention_dp=True"), + (dict(), dict(_TP16, moe_ep_size=16, moe_tp_size=1), None, "moe_ep=16"), + (dict(), dict(world_size=8, tp_size=8, moe_tp_size=8), None, "world_size=8"), + (dict(), _TP16, MTPDecodingConfig(max_draft_len=3, use_mtp_vanilla=True), "MTP"), + (dict(), _TP16, _NGRAM, "NGRAM"), + ], + ids=["other-shape", "attention-dp", "expert-parallel", "tp8", "vanilla-mtp", "ngram"], +) +def test_glm52_near_misses_do_not_match( + _on_sm107, config_kwargs, mapping_kwargs, spec_config, missed +): + config = _model_config( + _glm52_config(**config_kwargs), spec_config=spec_config, **mapping_kwargs + ) + assert modeling_v2_resolve(config, _llm_args("auto")) is None + with pytest.raises(ValueError, match=missed): + modeling_v2_resolve(config, _llm_args("require")) + + +def test_glm52_on_sm103_does_not_match(): + """The target's kernels are built for sm107; a GB300 falls back.""" + config = _model_config(_glm52_config(), **_TP16) + assert modeling_v2_resolve(config, _llm_args("auto")) is None + with pytest.raises(ValueError, match=r"compute capability \(10, 3\)"): + modeling_v2_resolve(config, _llm_args("require")) From e8989ffc21d1c0bd94d32bade71336c7f33bf3d3 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 15:55:35 -0700 Subject: [PATCH 06/53] Write the GLM-5.2 tp16 decode step out layer by layer DecodeTarget runs the embedding, the 78 decoder layers and the final norm as its own code instead of calling the built-in model's forward: the dense layers' attention all-reduce, residual norm, MLP and fused post-MLP all-reduce, and the MoE layers' fused post-attention all-reduce, experts and fused post-MoE all-reduce. The built-in modules still compute each part. The configuration the step is written for (fusions, MLP TP, multi-node experts, unquantized kv_a) is checked once when the target is built. Signed-off-by: Dongfeng Yu --- .../glm_5_2_nvfp4__sm_107__tp16/modeling.py | 106 +++++++++++++++++- 1 file changed, 100 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py index e7e86a427b58..38d3e1d3d766 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py @@ -20,9 +20,11 @@ from transformers import PretrainedConfig from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata +from tensorrt_llm._torch.distributed import AllReduceFusionOp, AllReduceParams from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3ForCausalLM, DeepseekV3Model from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP from tensorrt_llm._torch.speculative import SpecMetadata from . import weights as _weights @@ -64,10 +66,38 @@ class DecodeTarget: Routed to only when the batch has no context rows, which the core decides per step; nothing in this class reads the phase back. + + The step is written out layer by layer for this deployment's configuration, + which ``__init__`` checks once: + + - Layers 0-2 (dense, unquantized): the attention all-reduces its own output, + the residual add and ``post_attention_layernorm`` follow, and the MLP's + all-reduce carries the residual add and the next layer's input norm. + - Layers 3-77 (MoE): both all-reduces carry the residual add and the + following norm; the experts finalize before theirs (16 ranks span nodes). + - The last all-reduce applies the model's final norm (``next_layer_layernorm``). """ def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: self.core = core + for layer in core.layers[: core.num_hidden_layers]: + assert layer.allreduce is not None and layer.next_layer_layernorm is not None + assert layer.input_layernorm.nvfp4_scale is None, "kv_a_proj is not NVFP4" + if isinstance(layer.mlp, GatedMLP): + assert ( + layer.mlp_tp_size == _TP + and not layer.disable_attn_allreduce + and not layer.fusion_config.PRE_MLP_FUSION + and layer.fusion_config.POST_MLP_FUSION + and layer.post_attention_layernorm.nvfp4_scale is None + ), f"layer {layer.layer_idx}: unexpected dense-layer configuration" + else: + assert ( + layer.disable_attn_allreduce + and layer.fusion_config.PRE_MOE_FUSION + and layer.fusion_config.POST_MOE_FUSION + and layer.mapping.is_multi_node() + ), f"layer {layer.layer_idx}: unexpected MoE-layer configuration" def forward( self, @@ -78,14 +108,78 @@ def forward( spec_metadata: Optional[SpecMetadata] = None, **kwargs, ) -> torch.Tensor: - return DeepseekV3Model.forward( - self.core, + core = self.core + if inputs_embeds is None: + inputs_embeds = core.embed_tokens(input_ids) + residual = inputs_embeds + hidden_states = core.layers[0].input_layernorm(inputs_embeds) + for layer in core.layers[: core.num_hidden_layers]: + if isinstance(layer.mlp, GatedMLP): + hidden_states, residual = self._dense_layer( + layer, position_ids, hidden_states, residual, attn_metadata + ) + else: + hidden_states, residual = self._moe_layer( + layer, position_ids, hidden_states, residual, attn_metadata + ) + return hidden_states + + @staticmethod + def _dense_layer(layer, position_ids, hidden_states, residual, attn_metadata): + hidden_states = layer.self_attn( + position_ids=position_ids, + hidden_states=hidden_states, attn_metadata=attn_metadata, - input_ids=input_ids, + all_reduce_params=AllReduceParams(enable_allreduce=True), + ) + hidden_states, residual = layer.post_attention_layernorm(hidden_states, residual) + hidden_states = layer.mlp( + hidden_states, final_all_reduce_params=AllReduceParams(enable_allreduce=False) + ) + return layer.allreduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=layer.next_layer_layernorm.weight, + eps=layer.next_layer_layernorm.variance_epsilon, + ), + ) + + @staticmethod + def _moe_layer(layer, position_ids, hidden_states, residual, attn_metadata): + hidden_states = layer.self_attn( position_ids=position_ids, - inputs_embeds=inputs_embeds, - spec_metadata=spec_metadata, - **kwargs, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + all_reduce_params=AllReduceParams(enable_allreduce=False), + ) + hidden_states, residual = layer.allreduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=layer.post_attention_layernorm.weight, + eps=layer.post_attention_layernorm.variance_epsilon, + trigger_completion_at_end=False, + ), + ) + hidden_states = layer.mlp( + hidden_states, + None, + all_rank_num_tokens=attn_metadata.all_rank_num_tokens, + final_all_reduce_params=AllReduceParams(enable_allreduce=False), + do_finalize=True, + ) + return layer.allreduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=layer.next_layer_layernorm.weight, + eps=layer.next_layer_layernorm.variance_epsilon, + trigger_completion_at_end=False, + ), ) From 0045428682a1cac5ac9ed32687b7ba0cdf828a68 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 15:50:33 -0700 Subject: [PATCH 07/53] Add the GLM-5.2 TP16 small-M dense MLP op torch.ops.trtllm.glm_dense_mlp_small_m runs one rank's slice of a BF16 SwiGLU dense MLP for 1-8 rows as two weight-streaming kernels (kernels/glmMoeDsa/glmDenseMlpSmallM.cu): gate/up + SwiGLU on 192 CTAs in two-CTA clusters that split the 6144 contraction, then down on 96 CTAs of 64 output columns. The compiled tilings serve hidden 6144 with a rank-local intermediate size of 768, GLM-5.2's dense layers 0-2 at TP16. The op takes x [M, 6144], gate_up_weight [2 * 768, 6144] (this rank's gate rows over its up rows) and down_weight [6144, 768], and returns the rank-local partial [M, 6144] that the post-MLP all-reduce sums. The gate and up sums round to BF16 before SwiGLU and h rounds to BF16 before down, the rounding points of the unfused BF16 chain; the sums are FP32 in a fixed order, so repeats are bitwise equal. The sigmoid takes the fast exp and division. Under programmatic dependent launch every CTA bulk-copies its weights into shared memory before griddepcontrol.wait and reads its activations only after it. gate/up releases down at entry; down releases its dependents once its weight copies are issued. torch.ops.trtllm.glm_dense_mlp_small_m_supported(rows, hidden, inter) answers on the host whether the kernels serve a shape; the op raises on any other shape, layout, dtype or a device below sm_90. The target's entry is ops/gemm/dense_mlp_small_m.py, certified on sm_107 by its test against an FP32-sum reference with the same rounding points. Signed-off-by: Dongfeng Yu --- .../kernels/glmMoeDsa/glmDenseMlpSmallM.cu | 609 ++++++++++++++++++ .../kernels/glmMoeDsa/glmDenseMlpSmallM.h | 62 ++ cpp/tensorrt_llm/thop/CMakeLists.txt | 1 + cpp/tensorrt_llm/thop/glmDenseMlpSmallMOp.cpp | 115 ++++ .../ops/gemm/dense_mlp_small_m.py | 131 ++++ .../_torch/custom_ops/cpp_custom_ops.py | 4 + .../test_modeling_v2_dense_mlp_small_m.py | 246 +++++++ 7 files changed, 1168 insertions(+) create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.cu create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.h create mode 100644 cpp/tensorrt_llm/thop/glmDenseMlpSmallMOp.cpp create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/dense_mlp_small_m.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_dense_mlp_small_m.py diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.cu new file mode 100644 index 000000000000..b2178214c4ae --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.cu @@ -0,0 +1,609 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Weight-streaming BF16 kernels for the dense (non-MoE) MLP at decode row counts (M <= 8): the gate/up projection +// with SwiGLU, and the down projection. The compiled tilings serve hidden 6144 with a rank-local intermediate size of +// 768, the dense layers 0-2 of GLM-5.2 at TP16; the result is the rank-local partial that the post-MLP all-reduce sums. +// +// Each CTA owns 16-row weight tiles and its warps split the contraction. For every 32-wide contraction step a thread +// takes 16 bytes of two weight rows per tile and of one activation row and feeds them to two m16n8k16 MMAs, the +// weight on the MMA M side and the (at most 8) activation rows on the N side. The contraction index of an MMA k-slot +// depends only on the thread's position inside its quad, so the weight and activation fragments use the same +// permutation of k: the result is the same sum in another FP32 order. The gate/up kernel holds the gate rows and the +// up rows of the same output columns, and the thread that finishes a column applies SwiGLU in registers. +// +// Weights never depend on the preceding kernel: before griddepcontrol.wait each CTA copies all of its weights into +// shared memory with one bulk copy per weight row. After the wait each thread loads its activations from L2 and runs +// its MMAs; the warp sums are added in warp order in shared memory. With kClusterK > 1 the CTAs of a cluster split +// the contraction: each sends its sums to the first CTA of the cluster with st.async, which adds them in rank order. +// The results do not depend on timing. +// +// Programmatic dependent launch: activations are read only after griddepcontrol.wait; the dependents are released at +// entry or after the weight copies are issued (Trigger). Each grid fits in one wave of resident CTAs. + +#include "tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::dense_mlp_small_m +{ + +namespace +{ + +using Bf16 = __nv_bfloat16; + +constexpr int kWarpSize = 32; +//! Weight rows of one MMA tile (MMA M). +constexpr int kTileRows = 16; +//! BF16 values per 16-byte vector load. +constexpr int kVecElems = 8; +//! Threads per MMA quad; a quad covers kQuadThreads * kVecElems contraction elements per step. +constexpr int kQuadThreads = 4; +//! Contraction elements consumed by one warp step (two m16n8k16 MMAs per row tile). +constexpr int kKPerStep = kQuadThreads * kVecElems; +constexpr int kStepBytes = kKPerStep * static_cast(sizeof(Bf16)); + +//! When the dependents of a grid are released. +enum class Trigger : int +{ + kAtEntry, //!< first instruction + kAfterIssue, //!< after the weight copies are issued, before griddepcontrol.wait +}; + +//! Compile-time tiling of one kernel instance. +//! \tparam kGateUp gate/up + SwiGLU (true) or down projection (false). +//! \tparam kRowTiles 16-row weight tiles per CTA. gate/up: one tile holds 8 gate and the matching 8 up rows; with +//! more tiles, the first half are gate rows and the second half the up rows of the same columns. +//! \tparam kWarps warps per CTA; each warp owns a contiguous slice of the CTA's contraction. +//! \tparam kSteps 32-wide contraction steps per warp. +//! \tparam kClusterK CTAs of a cluster that split the contraction. +template +struct Tiling +{ + static constexpr bool kGateUp = kGateUp_; + static constexpr int kRowTiles = kRowTiles_; + static constexpr int kWarps = kWarps_; + static constexpr int kSteps = kSteps_; + static constexpr int kClusterK = kClusterK_; + static constexpr Trigger kTrigger = kTrigger_; + + static constexpr int kThreads = kWarps * kWarpSize; + static constexpr int kRows = kRowTiles * kTileRows; + static constexpr int kKPerWarp = kSteps * kKPerStep; + static constexpr int kKPerCta = kWarps * kKPerWarp; + //! The contraction size this tiling serves. + static constexpr int kKDim = kKPerCta * kClusterK; + //! Output columns of a CTA (gate/up: of h, down: of out). + static constexpr int kOutCols = kGateUp ? kRows / 2 : kRows; + //! (column, token) values the cluster finishes; each holds kEpiVals sums (gate and up, or one). + static constexpr int kEpi = kOutCols * kMaxRows; + static constexpr int kEpiVals = kGateUp ? 2 : 1; + + //! Weight stage [row][kKPerCta]: the bytes of one row copy, rows padded to 64 mod 128 bytes so that the 16-byte + //! fragment loads of a quarter warp (two rows, four quads) hit distinct banks. + static constexpr int kRowBytes = kKPerCta * static_cast(sizeof(Bf16)); + static constexpr int kRowStride = kRowBytes % 128 == 0 ? kRowBytes + 64 : kRowBytes; + + // Shared-memory layout (bytes): warp sums | sums received by the cluster leader | weight stage | barriers. + static constexpr int kRedBytes = kWarps * kRowTiles * 4 * kWarpSize * static_cast(sizeof(float)); + static constexpr int kRecvBytes = (kClusterK - 1) * kEpi * kEpiVals * static_cast(sizeof(float)); + static constexpr int kStageOffset = (kRedBytes + kRecvBytes + 127) / 128 * 128; + static constexpr int kBarOffset = kStageOffset + kRows * kRowStride; + static constexpr int kSmemBytes = kBarOffset + 2 * static_cast(sizeof(uint64_t)); + + static_assert(!kGateUp_ || kRowTiles_ == 1 || kRowTiles_ % 2 == 0, "gate/up needs one or an even number of tiles"); + static_assert(kRows <= kThreads, "one thread issues the copy of each weight row"); + static_assert(kRowBytes % 64 == 0, "rows are whole contraction steps"); +}; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + +__device__ __forceinline__ void pdlWait() +{ + asm volatile("griddepcontrol.wait;" ::: "memory"); +} + +__device__ __forceinline__ void pdlLaunchDependents() +{ + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +} + +//! Activation load from L2 (the preceding kernel wrote it; L1 is bypassed). +__device__ __forceinline__ uint4 loadActivation(Bf16 const* ptr) +{ + uint4 r; + asm volatile("ld.global.cg.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w) + : "l"(ptr) + : "memory"); + return r; +} + +__device__ __forceinline__ uint32_t smemAddress(void const* ptr) +{ + return static_cast(__cvta_generic_to_shared(ptr)); +} + +//! One arrival (this one) plus expectBytes of transactions. +__device__ __forceinline__ void mbarrierInit(uint32_t bar, uint32_t expectBytes) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(bar) : "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(bar), "r"(expectBytes) : "memory"); +} + +__device__ __forceinline__ void mbarrierWait(uint32_t bar) +{ + asm volatile( + "{\n" + ".reg .pred done;\n" + "WAIT_%=:\n" + "mbarrier.try_wait.parity.shared::cta.b64 done, [%0], 0;\n" + "@!done bra WAIT_%=;\n" + "}\n" ::"r"(bar) + : "memory"); +} + +__device__ __forceinline__ void mbarrierWaitCluster(uint32_t bar) +{ + asm volatile( + "{\n" + ".reg .pred done;\n" + "WAIT_%=:\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 done, [%0], 0;\n" + "@!done bra WAIT_%=;\n" + "}\n" ::"r"(bar) + : "memory"); +} + +__device__ __forceinline__ void bulkCopy(uint32_t dst, void const* src, uint32_t bytes, uint32_t bar) +{ + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"(dst), + "l"(src), "r"(bytes), "r"(bar) + : "memory"); +} + +__device__ __forceinline__ uint32_t clusterCtaRank() +{ + uint32_t rank; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(rank)); + return rank; +} + +__device__ __forceinline__ uint32_t mapToCta(uint32_t addr, uint32_t rank) +{ + uint32_t peer; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(peer) : "r"(addr), "r"(rank)); + return peer; +} + +__device__ __forceinline__ void mmaBf16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) +{ + asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +//! Two MMAs of one step and row tile: MMA A takes contraction elements 0..3 of the thread's 8-element vectors, MMA B +//! elements 4..7. Fragments: a0/a2 = weight row "group", a1/a3 = weight row "group + 8", b0/b1 = activation row +//! "group". +__device__ __forceinline__ void mmaStep(float (&accA)[4], float (&accB)[4], uint4 const& wLo, uint4 const& wHi, + uint4 const& xv) +{ + mmaBf16(accA, wLo.x, wHi.x, wLo.y, wHi.y, xv.x, xv.y); + mmaBf16(accB, wLo.z, wHi.z, wLo.w, wHi.w, xv.z, xv.w); +} + +__device__ __forceinline__ float roundToBf16(float v) +{ + return __bfloat162float(__float2bfloat16_rn(v)); +} + +//! SwiGLU of BF16-rounded gate and up sums, as the BF16 GEMM output feeding silu_and_mul: the sigmoid takes the fast +//! exp and division of that kernel. +__device__ __forceinline__ float swiglu(float gateSum, float upSum) +{ + float const gate = roundToBf16(gateSum); + float const up = roundToBf16(upSum); + float const sig = __fdividef(1.F, 1.F + __expf(-gate)); + return gate * sig * up; +} + +//! Weight row (of the [*, kDim] weight matrix) of row r of row tile t of the CTA of output tile `tile`. +template +__device__ __forceinline__ int weightRow(int tile, int t, int r, int inter) +{ + if constexpr (T::kGateUp) + { + if constexpr (T::kRowTiles == 1) + { + int const col = tile * T::kOutCols + (r & 7); + return r < 8 ? col : inter + col; + } + else + { + constexpr int kHalf = T::kRowTiles / 2; + int const col = tile * T::kOutCols + (t % kHalf) * kTileRows + r; + return t < kHalf ? col : inter + col; + } + } + else + { + return tile * T::kRows + t * kTileRows + r; + } +} + +#endif // __CUDA_ARCH__ >= 900 + +//! x [numRows, kDim], weight [*, kDim]; out: gate/up h [numRows, outDim] (outDim = inter), down out [numRows, outDim]. +template +__global__ void __launch_bounds__(T::kThreads, 1) denseMlpSmallMKernel(Bf16 const* __restrict__ x, + Bf16 const* __restrict__ weight, Bf16* __restrict__ out, int numRows, int outDim) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + constexpr int kDim = T::kKDim; + extern __shared__ __align__(128) char smem[]; + float* red = reinterpret_cast(smem); + [[maybe_unused]] float* recv = reinterpret_cast(smem + T::kRedBytes); + char* stage = smem + T::kStageOffset; + uint32_t const stageBar = smemAddress(smem + T::kBarOffset); + [[maybe_unused]] uint32_t const clusterBar = stageBar + static_cast(sizeof(uint64_t)); + + if constexpr (T::kTrigger == Trigger::kAtEntry) + { + pdlLaunchDependents(); + } + + int const lane = threadIdx.x % kWarpSize; + int const warp = threadIdx.x / kWarpSize; + int const group = lane / kQuadThreads; + int const tig = lane % kQuadThreads; + int const rank = T::kClusterK > 1 ? static_cast(clusterCtaRank()) : 0; + int const tile = blockIdx.x / T::kClusterK; + int const kCta = rank * T::kKPerCta; + + // The weight stage expects all row bytes; the cluster leader also expects the sums of the other CTAs. + if (threadIdx.x == 0) + { + mbarrierInit(stageBar, static_cast(T::kRows * T::kRowBytes)); + if constexpr (T::kClusterK > 1) + { + if (rank == 0) + { + mbarrierInit(clusterBar, static_cast(T::kRecvBytes)); + } + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + __syncthreads(); + if constexpr (T::kClusterK > 1) + { + // Publishes the leader's barrier; every CTA waits for this phase before it sends. + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); + } + + // Thread r copies weight row r of the CTA (its contraction slice). + int const copyRow = static_cast(threadIdx.x); + if (copyRow < T::kRows) + { + int const wRow = weightRow(tile, copyRow / kTileRows, copyRow % kTileRows, outDim); + bulkCopy(smemAddress(stage + copyRow * T::kRowStride), weight + static_cast(wRow) * kDim + kCta, + static_cast(T::kRowBytes), stageBar); + } + + if constexpr (T::kTrigger == Trigger::kAfterIssue) + { + pdlLaunchDependents(); + } + pdlWait(); + + // Activations: row "group" (zero beyond numRows), the thread's 8 contraction elements of every step. + int const kWarp = kCta + warp * T::kKPerWarp; + bool const hasRow = group < numRows; + Bf16 const* xRow = x + static_cast(hasRow ? group : 0) * kDim + kWarp + tig * kVecElems; + uint4 xv[T::kSteps]; +#pragma unroll + for (int s = 0; s < T::kSteps; ++s) + { + xv[s] = hasRow ? loadActivation(xRow + s * kKPerStep) : make_uint4(0U, 0U, 0U, 0U); + } + + float accA[T::kRowTiles][4]; + float accB[T::kRowTiles][4]; +#pragma unroll + for (int t = 0; t < T::kRowTiles; ++t) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + accA[t][e] = 0.F; + accB[t][e] = 0.F; + } + } + mbarrierWait(stageBar); + char const* seg = stage + warp * T::kKPerWarp * static_cast(sizeof(Bf16)) + tig * 16; +#pragma unroll + for (int s = 0; s < T::kSteps; ++s) + { +#pragma unroll + for (int t = 0; t < T::kRowTiles; ++t) + { + int const rowLo = t * kTileRows + group; + uint4 const wLo = *reinterpret_cast(seg + rowLo * T::kRowStride + s * kStepBytes); + uint4 const wHi = *reinterpret_cast(seg + (rowLo + 8) * T::kRowStride + s * kStepBytes); + mmaStep(accA[t], accB[t], wLo, wHi, xv[s]); + } + } + + // Warp sums -> shared memory [warp][tile * 4 + element][lane]. +#pragma unroll + for (int t = 0; t < T::kRowTiles; ++t) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + red[(warp * T::kRowTiles * 4 + t * 4 + e) * kWarpSize + lane] = accA[t][e] + accB[t][e]; + } + } + __syncthreads(); + if constexpr (T::kClusterK > 1) + { + asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); + } + + // Finish value j: element (tile t, accumulator e, lane l) holds weight row t * 16 + l / 4 + (e / 2) * 8 of the + // CTA and token 2 * (l % 4) + e % 2. + for (int j = static_cast(threadIdx.x); j < T::kEpi; j += T::kThreads) + { + int q0; + [[maybe_unused]] int q1 = 0; + int col; + int token; + int const l = j % kWarpSize; + if constexpr (T::kGateUp && T::kRowTiles == 1) + { + int const e = j / kWarpSize; // 0 or 1: rows l / 4 (gate) and l / 4 + 8 (up) + q0 = e; + q1 = e + 2; + col = l / kQuadThreads; + token = 2 * (l % kQuadThreads) + e; + } + else if constexpr (T::kGateUp) + { + int const t = j / (4 * kWarpSize); + int const e = (j / kWarpSize) % 4; + q0 = t * 4 + e; + q1 = (t + T::kRowTiles / 2) * 4 + e; + col = t * kTileRows + l / kQuadThreads + (e / 2) * 8; + token = 2 * (l % kQuadThreads) + e % 2; + } + else + { + int const q = j / kWarpSize; + q0 = q; + col = (q / 4) * kTileRows + l / kQuadThreads + ((q % 4) / 2) * 8; + token = 2 * (l % kQuadThreads) + q % 2; + } + // Fixed order over warps keeps the result deterministic. + float v0 = 0.F; + [[maybe_unused]] float v1 = 0.F; +#pragma unroll + for (int w = 0; w < T::kWarps; ++w) + { + v0 += red[(w * T::kRowTiles * 4 + q0) * kWarpSize + l]; + if constexpr (T::kGateUp) + { + v1 += red[(w * T::kRowTiles * 4 + q1) * kWarpSize + l]; + } + } + if constexpr (T::kClusterK > 1) + { + if (rank != 0) + { + uint32_t const slot = smemAddress(recv + ((rank - 1) * T::kEpi + j) * T::kEpiVals); + uint32_t const peerSlot = mapToCta(slot, 0); + uint32_t const peerBar = mapToCta(clusterBar, 0); + if constexpr (T::kGateUp) + { + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.b32 [%0], {%1, %2}, [%3];" ::"r"( + peerSlot), + "r"(__float_as_uint(v0)), "r"(__float_as_uint(v1)), "r"(peerBar) + : "memory"); + } + else + { + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"( + peerSlot), + "r"(__float_as_uint(v0)), "r"(peerBar) + : "memory"); + } + continue; + } + mbarrierWaitCluster(clusterBar); + // Rank order. +#pragma unroll + for (int r = 1; r < T::kClusterK; ++r) + { + v0 += recv[((r - 1) * T::kEpi + j) * T::kEpiVals]; + if constexpr (T::kGateUp) + { + v1 += recv[((r - 1) * T::kEpi + j) * T::kEpiVals + 1]; + } + } + } + if (token < numRows) + { + int64_t const idx = static_cast(token) * outDim + tile * T::kOutCols + col; + if constexpr (T::kGateUp) + { + out[idx] = __float2bfloat16_rn(swiglu(v0, v1)); + } + else + { + out[idx] = __float2bfloat16_rn(v0); + } + } + } +#endif +} + +template +cudaError_t launchTiling(Bf16 const* in, Bf16 const* weight, Bf16* out, int numRows, int outDim, bool enablePdl, + cudaStream_t stream) +{ + auto const kernel = &denseMlpSmallMKernel; + if constexpr (T::kSmemBytes > 48 * 1024) + { + static cudaError_t const attrStatus + = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, T::kSmemBytes); + if (attrStatus != cudaSuccess) + { + return attrStatus; + } + } + cudaLaunchConfig_t config{}; + config.gridDim = dim3(outDim / T::kOutCols * T::kClusterK); + config.blockDim = dim3(T::kThreads); + config.dynamicSmemBytes = T::kSmemBytes; + config.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enablePdl ? 1 : 0; + attrs[1].id = cudaLaunchAttributeClusterDimension; + attrs[1].val.clusterDim.x = T::kClusterK; + attrs[1].val.clusterDim.y = 1; + attrs[1].val.clusterDim.z = 1; + config.attrs = attrs; + config.numAttrs = T::kClusterK > 1 ? 2 : 1; + return cudaLaunchKernelEx(&config, kernel, in, weight, out, numRows, outDim); +} + +// Tilings: row tiles, warps, steps, cluster size, trigger. +using GateUpTilings = std::tuple< + // Hidden 6144 (GLM-5.2): 2-CTA clusters split the contraction, 16 warps x 6 steps, 8 columns per cluster; at TP16 + // (inter 768) 192 CTAs with 96 KiB of weights each. The down kernel is released at entry, so its weights are staged + // during the preceding all-reduce as well. + Tiling>; + +using DownTilings = std::tuple< + // Inter 768 (GLM-5.2 at TP16): 64 rows per CTA, 8 warps x 3 steps; at hidden 6144 96 CTAs with 96 KiB of weights + // each. + Tiling>; + +template +constexpr int kNumTilings = static_cast(std::tuple_size_v); + +template +bool tilingSupported(int tiling, int numRows, int kDim, int outDim) +{ + if constexpr (I < kNumTilings) + { + if (tiling == I) + { + using T = std::tuple_element_t; + return numRows >= 1 && numRows <= kMaxRows && kDim == T::kKDim && outDim > 0 + && outDim % T::kOutCols == 0; + } + return tilingSupported(tiling, numRows, kDim, outDim); + } + else + { + return false; + } +} + +template +cudaError_t launchTilingIndex(int tiling, Bf16 const* in, Bf16 const* weight, Bf16* out, int numRows, int outDim, + bool enablePdl, cudaStream_t stream) +{ + if constexpr (I < kNumTilings) + { + if (tiling == I) + { + return launchTiling>( + in, weight, out, numRows, outDim, enablePdl, stream); + } + return launchTilingIndex(tiling, in, weight, out, numRows, outDim, enablePdl, stream); + } + else + { + return cudaErrorInvalidValue; + } +} + +//! The tiling a call runs: the requested one if it serves the shape, for kDefaultTiling the first entry that does; +//! -1 if none. +template +int resolveTiling(int tiling, int numRows, int kDim, int outDim) +{ + if (tiling != kDefaultTiling) + { + return tilingSupported(tiling, numRows, kDim, outDim) ? tiling : -1; + } + for (int i = 0; i < kNumTilings; ++i) + { + if (tilingSupported(i, numRows, kDim, outDim)) + { + return i; + } + } + return -1; +} + +} // namespace + +bool isGateUpSupported(int numRows, int hidden, int inter, int tiling) +{ + return resolveTiling(tiling, numRows, hidden, inter) >= 0; +} + +bool isDownSupported(int numRows, int inter, int hidden, int tiling) +{ + return resolveTiling(tiling, numRows, inter, hidden) >= 0; +} + +cudaError_t launchGateUpSwiglu(Bf16 const* x, Bf16 const* gateUpWeight, Bf16* h, int numRows, int hidden, int inter, + int tiling, bool enablePdl, cudaStream_t stream) +{ + int const resolved = resolveTiling(tiling, numRows, hidden, inter); + if (resolved < 0) + { + return cudaErrorInvalidValue; + } + return launchTilingIndex(resolved, x, gateUpWeight, h, numRows, inter, enablePdl, stream); +} + +cudaError_t launchDown(Bf16 const* h, Bf16 const* downWeight, Bf16* out, int numRows, int inter, int hidden, + int tiling, bool enablePdl, cudaStream_t stream) +{ + int const resolved = resolveTiling(tiling, numRows, inter, hidden); + if (resolved < 0) + { + return cudaErrorInvalidValue; + } + return launchTilingIndex(resolved, h, downWeight, out, numRows, hidden, enablePdl, stream); +} + +} // namespace kernels::glm_moe_dsa::dense_mlp_small_m + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.h new file mode 100644 index 000000000000..90cd1b7f90a6 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.h @@ -0,0 +1,62 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::dense_mlp_small_m +{ + +//! Largest number of activation rows (tokens) the kernels accept. +constexpr int kMaxRows = 8; + +//! Tiling argument that selects the default tiling of the shape. +constexpr int kDefaultTiling = -1; + +//! Host-side eligibility checks (no device work). tiling = kDefaultTiling checks the default tiling of the shape. +//! \param numRows activation rows M. +//! \param hidden model hidden size (gate/up contraction size, down output size). +//! \param inter rank-local intermediate size (gate/up output size, down contraction size). +bool isGateUpSupported(int numRows, int hidden, int inter, int tiling = kDefaultTiling); +bool isDownSupported(int numRows, int inter, int hidden, int tiling = kDefaultTiling); + +//! Gate/up projection fused with SwiGLU. For every row m < numRows and column j < inter: +//! gate = bf16(sum_k x[m, k] * W[j, k]), up = bf16(sum_k x[m, k] * W[inter + j, k]), +//! h[m, j] = bf16((gate * sigmoid(gate)) * up), sums in FP32, sigmoid as 1 / (1 + exp(-gate)) with the fast exp +//! and division. +//! \param x activations [numRows, hidden], row-major, BF16, 16-byte aligned. +//! \param gateUpWeight [2 * inter, hidden], row-major, BF16 (rows [0, inter) gate, [inter, 2 * inter) up). +//! \param h output [numRows, inter], row-major, BF16. +cudaError_t launchGateUpSwiglu(__nv_bfloat16 const* x, __nv_bfloat16 const* gateUpWeight, __nv_bfloat16* h, + int numRows, int hidden, int inter, int tiling, bool enablePdl, cudaStream_t stream); + +//! Down projection: out[m, n] = bf16(sum_k h[m, k] * W[n, k]) with FP32 sums. +//! \param h activations [numRows, inter], row-major, BF16, 16-byte aligned. +//! \param downWeight [hidden, inter], row-major, BF16. +//! \param out output [numRows, hidden], row-major, BF16. +cudaError_t launchDown(__nv_bfloat16 const* h, __nv_bfloat16 const* downWeight, __nv_bfloat16* out, int numRows, + int inter, int hidden, int tiling, bool enablePdl, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa::dense_mlp_small_m + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index ab424318c862..ce6ad1d3eae7 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -98,6 +98,7 @@ add_th_common_family( fp8RowwiseGemm.cpp fp8Quantize.cpp dsv3FusedAGemmOp.cpp + glmDenseMlpSmallMOp.cpp deepseekV4QNormOp.cpp fusedAddRMSNormQuant.cpp fusedActivationQuant.cpp diff --git a/cpp/tensorrt_llm/thop/glmDenseMlpSmallMOp.cpp b/cpp/tensorrt_llm/thop/glmDenseMlpSmallMOp.cpp new file mode 100644 index 000000000000..91b36a423ec4 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmDenseMlpSmallMOp.cpp @@ -0,0 +1,115 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace tkd = tensorrt_llm::kernels::glm_moe_dsa::dense_mlp_small_m; + +//! The kernels use clusters, bulk copies and st.async; below sm_90 their bodies compile to nothing. +constexpr int kMinSmVersion = 90; + +void checkBf16Matrix(at::Tensor const& t, char const* name) +{ + CHECK_INPUT(t, at::ScalarType::BFloat16); + TORCH_CHECK(t.dim() == 2, name, " must be 2-D"); + TORCH_CHECK(reinterpret_cast(t.data_ptr()) % 16 == 0, name, " must be 16-byte aligned"); +} + +} // namespace + +//! Whether the kernels serve M rows of a dense MLP with this hidden and rank-local intermediate size (host only). +bool glmDenseMlpSmallMSupported(int64_t numRows, int64_t hidden, int64_t inter) +{ + return tkd::isGateUpSupported(static_cast(numRows), static_cast(hidden), static_cast(inter)) + && tkd::isDownSupported(static_cast(numRows), static_cast(inter), static_cast(hidden)); +} + +//! The rank-local partial of a BF16 gated MLP with plain SwiGLU (no output all-reduce): x [M, H], +//! gate_up_weight [2 I, H] (gate rows first), down_weight [H, I] -> [M, H]. gate_up_tiling / down_tiling select a +//! compiled tiling (-1: the default of the shape). +at::Tensor glmDenseMlpSmallM(at::Tensor const& x, at::Tensor const& gateUpWeight, at::Tensor const& downWeight, + int64_t gateUpTiling, int64_t downTiling) +{ + checkBf16Matrix(x, "x"); + checkBf16Matrix(gateUpWeight, "gate_up_weight"); + checkBf16Matrix(downWeight, "down_weight"); + TORCH_CHECK(gateUpWeight.device() == x.device() && downWeight.device() == x.device(), + "glm_dense_mlp_small_m: all tensors must be on the device of x"); + TORCH_CHECK(tensorrt_llm::common::getSMVersion() >= kMinSmVersion, "glm_dense_mlp_small_m needs sm_90 or newer"); + int64_t const numRows = x.size(0); + int64_t const hidden = x.size(1); + int64_t const inter = downWeight.size(1); + TORCH_CHECK(gateUpWeight.size(0) == 2 * inter && gateUpWeight.size(1) == hidden && downWeight.size(0) == hidden, + "glm_dense_mlp_small_m: gate_up_weight ", gateUpWeight.sizes(), " and down_weight ", downWeight.sizes(), + " do not match x ", x.sizes()); + TORCH_CHECK(tkd::isGateUpSupported(static_cast(numRows), static_cast(hidden), static_cast(inter), + static_cast(gateUpTiling)) + && tkd::isDownSupported(static_cast(numRows), static_cast(inter), static_cast(hidden), + static_cast(downTiling)), + "glm_dense_mlp_small_m: unsupported shape (M ", numRows, ", hidden ", hidden, ", inter ", inter, ") or tiling"); + + at::Tensor h = at::empty({numRows, inter}, x.options()); + at::Tensor out = at::empty({numRows, hidden}, x.options()); + auto const stream = at::cuda::getCurrentCUDAStream(x.get_device()); + bool const enablePdl = tensorrt_llm::common::getEnvEnablePDL(); + TLLM_CUDA_CHECK(tkd::launchGateUpSwiglu(static_cast<__nv_bfloat16 const*>(x.data_ptr()), + static_cast<__nv_bfloat16 const*>(gateUpWeight.data_ptr()), static_cast<__nv_bfloat16*>(h.data_ptr()), + static_cast(numRows), static_cast(hidden), static_cast(inter), static_cast(gateUpTiling), + enablePdl, stream)); + TLLM_CUDA_CHECK(tkd::launchDown(static_cast<__nv_bfloat16 const*>(h.data_ptr()), + static_cast<__nv_bfloat16 const*>(downWeight.data_ptr()), static_cast<__nv_bfloat16*>(out.data_ptr()), + static_cast(numRows), static_cast(inter), static_cast(hidden), static_cast(downTiling), + enablePdl, stream)); + return out; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("glm_dense_mlp_small_m_supported(int num_rows, int hidden, int inter) -> bool"); + m.def( + "glm_dense_mlp_small_m(Tensor x, Tensor gate_up_weight, Tensor down_weight, int gate_up_tiling=-1, " + "int down_tiling=-1) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CompositeExplicitAutograd, m) +{ + m.impl("glm_dense_mlp_small_m_supported", &tensorrt_llm::torch_ext::glmDenseMlpSmallMSupported); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_dense_mlp_small_m", &tensorrt_llm::torch_ext::glmDenseMlpSmallM); +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/dense_mlp_small_m.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/dense_mlp_small_m.py new file mode 100644 index 000000000000..577c93008ffd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/dense_mlp_small_m.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One rank's slice of the BF16 dense MLP for up to 8 decode rows, via the trtllm +glm_dense_mlp_small_m op (two weight-streaming kernels).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +#: GLM-5.2's hidden size, the contraction of gate/up and the output width of down. +HIDDEN = 6144 +#: The rank-local intermediate size of the dense layers 0-2 at TP16: 12288 / 16. +INTER_PER_RANK = 768 + + +class GlmDenseMlpSmallM(OpWrapper): + """`out = (silu(x @ W_gate.T) * (x @ W_up.T)) @ W_down.T` on one rank's slice, 1-8 rows. + + The slice is TP16's: `gate_up_weight` [2 * 768, 6144] stacks this rank's + 768 gate_proj rows over the same 768 up_proj rows (gate first), and + `down_weight` [6144, 768] holds the same 768 input columns of down_proj. + `out` [M, 6144] is the rank-local partial: the sum over the 16 ranks, which + the post-MLP all-reduce forms, is the MLP output. Returns a new tensor; no + input is modified. + + Rounding points are the unfused chain's: the gate and up sums round to + bf16 (the bf16 gate_up GEMM output), SwiGLU runs in fp32 and rounds once to + bf16 `h`, and the down sum rounds to bf16. Sums are fp32 in another order. + + Two kernels on the current stream: gate/up + SwiGLU (192 CTAs in 2-CTA + clusters that split the 6144 contraction) and down (96 CTAs of 64 output + columns). Under programmatic dependent launch each CTA bulk-copies all of + its weights into shared memory before its grid-dependency wait, and reads + `x` (gate/up) or `h` (down) only after it. gate/up releases down at entry; + down releases its own dependents once its weight copies are issued, so the + kernel after the op must wait before it reads `out`. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The target calls the op on the dense layers 0-2 of a generation-only + # forward: rows = requests x (draft length + 1), at most 8. Operand scales + # are the checkpoint's: post_attention_layernorm's gamma averages 0.021 on + # these layers (x has unit RMS before it), gate/up weights 0.014 std, down + # weights 0.013 std. `saturated` scales gate/up so the gate sums reach the + # sigmoid's tails. + CELLS: tuple[Cell, ...] = ( + Cell( + why="c1 verify forward: MTP K=7's 8 rows through a TP16 dense layer", + spec=dict(rows=8, x_scale=0.021, gate_up_std=0.014, down_std=0.013), + ), + Cell( + why="decode without speculation: 1 row, 7 of the 8 MMA token columns empty", + spec=dict(rows=1, x_scale=0.021, gate_up_std=0.014, down_std=0.013), + ), + Cell( + why="MTP K=3: 4 rows, the empty token columns start at a thread's pair", + spec=dict(rows=4, x_scale=0.021, gate_up_std=0.014, down_std=0.013), + ), + Cell( + why="saturated SwiGLU: gate sums to +-40, the fast exp and division in sigmoid's tails", + spec=dict(rows=8, x_scale=1.0, gate_up_std=0.13, down_std=0.036), + ), + ) + + note = """ + Shapes: hidden 6144, rank-local intermediate 768, 1-8 rows; BF16, + contiguous, 16-byte aligned operands on an sm_90 or newer device. The op + raises on anything else, so the caller routes other shapes (more rows, + context forwards, another TP) to a generic gated MLP. + `torch.ops.trtllm.glm_dense_mlp_small_m_supported(rows, hidden, inter)` + answers the shape question on the host. + + The weight layout is the caller's to get right and no check can see it: a + gate_up_weight that stacks up over gate, or that takes rows + [r * 1536, (r + 1) * 1536) of a pre-fused [gate; up] tensor, gives a + plausible wrong MLP. + + The weights are read before the grid-dependency wait, so they must not be + written by any kernel still running ahead of the op on the stream. + """ + + def __call__( + self, x: torch.Tensor, gate_up_weight: torch.Tensor, down_weight: torch.Tensor + ) -> torch.Tensor: + return torch.ops.trtllm.glm_dense_mlp_small_m(x, gate_up_weight, down_weight) + + def reference( + self, x: torch.Tensor, gate_up_weight: torch.Tensor, down_weight: torch.Tensor + ) -> torch.Tensor: + """The rank partial with the op's rounding points, fp32 sums. + + The intermediate bf16 roundings are part of what the op computes (the + unfused bf16 GEMMs round there too), so the reference rounds at the + same points instead of carrying fp32 through; without them the band + would measure those roundings, not the kernel. + + TF32 off: torch's default fp32 matmul on this hardware is a TF32 + product, which would make the reference the inaccurate side. + """ + inter = down_weight.shape[1] + with true_fp32_matmul(): + gate_up = x.float() @ gate_up_weight.float().t() + gate = gate_up[:, :inter].to(torch.bfloat16).float() + up = gate_up[:, inter:].to(torch.bfloat16).float() + h = (torch.nn.functional.silu(gate) * up).to(torch.bfloat16) + with true_fp32_matmul(): + return (h.float() @ down_weight.float().t()).to(x.dtype) + + def is_valid( + self, x: torch.Tensor, gate_up_weight: torch.Tensor, down_weight: torch.Tensor + ) -> None: + """Nothing to refuse: every input the op accepts it computes as stated. + + The op raises on every shape, dtype, layout, alignment and device it + does not serve, including pre-sm_90 devices where the kernels have no + body. What is left is the weight layout (see `note`), which no check on + the tensors can see. + """ + return None + + +dense_mlp_small_m = GlmDenseMlpSmallM() diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index f500034c4df7..fd8961fa4711 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -985,6 +985,10 @@ def _(pe: torch.Tensor, nope: torch.Tensor): scale = pe.new_empty((M, 1), dtype=torch.int32) return packed, scale + @torch.library.register_fake("trtllm::glm_dense_mlp_small_m") + def _(x, gate_up_weight, down_weight, gate_up_tiling=-1, down_tiling=-1): + return x.new_empty((x.shape[0], down_weight.shape[0])) + @torch.library.register_fake("trtllm::fp8_quantize_1x128_packed_ue8m0") def _(input: torch.Tensor, use_r128c4_layout: bool = True) -> tuple[torch.Tensor, torch.Tensor]: diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_dense_mlp_small_m.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_dense_mlp_small_m.py new file mode 100644 index 000000000000..98184625ba61 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_dense_mlp_small_m.py @@ -0,0 +1,246 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM-5.2 tp16 dense_mlp_small_m entry (sm_107). + +The cells gate the kernels against the entry's reference within a derived band. +The rest is what a band cannot show: a row's result does not depend on the +other rows, repeats are bitwise equal, a call chained right behind another one +reads its input only after the first call wrote it (programmatic dependent +launch), CUDA-graph replay, and the shape domain the op serves and refuses. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.dense_mlp_small_m import ( # noqa: E501 + HIDDEN, + INTER_PER_RANK, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.dense_mlp_small_m import ( # noqa: E501 + dense_mlp_small_m as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +assert torch.cuda.is_available(), "dense_mlp_small_m requires a CUDA device" + +#: The c1 cell: the shape and scales every non-cell test below reuses. +C1 = op.CELLS[0] + + +def _seed(text: str) -> int: + return zlib.crc32(text.encode()) % 2**31 + + +def _build(spec, seed): + """x [rows, 6144], gate_up_weight [2 * 768, 6144] and down_weight [6144, 768], bf16. + + Normal draws at the cell's scales; x is a unit-RMS row times the norm's + gamma, as post_attention_layernorm hands it over. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + x = torch.randn(spec["rows"], HIDDEN, generator=g, device="cuda") * spec["x_scale"] + gate_up = torch.randn(2 * INTER_PER_RANK, HIDDEN, generator=g, device="cuda") + down = torch.randn(HIDDEN, INTER_PER_RANK, generator=g, device="cuda") + return ( + x.to(torch.bfloat16), + (gate_up * spec["gate_up_std"]).to(torch.bfloat16), + (down * spec["down_std"]).to(torch.bfloat16), + ) + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Summation order and the fast sigmoid, nothing else. + + Kernel and reference take the same bf16 operands and round at the same + three points (gate and up sums, h, out). They differ in the fp32 order of + the 6144- and 768-long sums (about 1e-6 relative) and in the kernel's fast + exp and division inside the sigmoid (a few fp32 ulp; tens at the saturated + cell's |gate| of 40). Either can carry a value across a bf16 rounding + boundary. An output that flips moves by one bf16 ulp of itself: at most + 2^-7 of its row's largest magnitude, 2 units of the 2^-8 band unit. An h + that flips moves each output by 2^-7 |h_k w_nk|, a few hundredths of a + unit. Flips are rare (the gaps above are 1e-3 of a bf16 ulp), so the + relative RMS sits well under one unit. + + So: 4 units per element against the row scale, 2 units of relative RMS -- + twice the single-flip bound, and three orders below a wrong product (a + swapped gate/up, a lost row, a missing SwiGLU), which misses by O(1). + """ + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + x, gate_up_weight, down_weight = _build(cell.spec, seed=_seed(cell.why)) + before = x.clone() + out = op(x, gate_up_weight, down_weight) + assert out.shape == (x.shape[0], HIDDEN) and out.dtype == torch.bfloat16 + assert out.is_contiguous() + assert torch.equal(x, before), "the op must not modify x" + _compare(out, op.reference(x, gate_up_weight, down_weight)) + + +def test_a_row_does_not_see_the_other_rows() -> None: + """Row m's result is bitwise the same with 1 to 8 rows present. + + The empty token columns of a short call read as zero and are never + stored, and the arithmetic of a row does not depend on the other rows. + """ + x, gate_up_weight, down_weight = _build(C1.spec, seed=_seed("rows")) + full = op(x, gate_up_weight, down_weight) + for rows in range(1, 8): + part = op(x[:rows].contiguous(), gate_up_weight, down_weight) + assert torch.equal(part, full[:rows]), f"{rows} rows differ from the 8-row call" + + +def test_repeats_are_bitwise_equal() -> None: + """Fixed summation orders inside a CTA, across warps and across the cluster.""" + x, gate_up_weight, down_weight = _build(C1.spec, seed=_seed("repeat")) + first = op(x, gate_up_weight, down_weight) + for _ in range(4): + assert torch.equal(op(x, gate_up_weight, down_weight), first) + + +def _chain(x, layer_a, layer_b): + out_a = op(x, *layer_a) + return out_a, op(out_a, *layer_b) + + +def test_a_chained_call_reads_its_input_after_the_wait() -> None: + """The second of two back-to-back calls reads what the first one wrote. + + With PDL the down kernel of the first call releases the gate/up kernel of + the second before it computes its output; that kernel stages its weights + early and must read x only after its grid-dependency wait. Every iteration + draws a new x, so a read of the previous iteration's output (often at the + same address) or of an unwritten buffer shows up against a standalone + recompute of the second call after a device sync. Eager and graph-replayed. + """ + x, *layer_a = _build(C1.spec, seed=_seed("chain a")) + _, *layer_b = _build(C1.spec, seed=_seed("chain b")) + g = torch.Generator(device="cuda").manual_seed(_seed("chain x")) + + for _ in range(8): + x.copy_(torch.randn(x.shape, generator=g, device="cuda") * C1.spec["x_scale"]) + out_a, out_b = _chain(x, layer_a, layer_b) + torch.cuda.synchronize() + assert torch.equal(out_b, op(out_a.clone(), *layer_b)), "eager chain read a stale input" + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + _chain(x, layer_a, layer_b) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out_a, out_b = _chain(x, layer_a, layer_b) + for _ in range(8): + x.copy_(torch.randn(x.shape, generator=g, device="cuda") * C1.spec["x_scale"]) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out_b, op(out_a.clone(), *layer_b)), "replayed chain read a stale input" + + +def test_cuda_graph_capture_and_replay() -> None: + """The c1 forward runs inside one CUDA graph: replay with new x contents + gives the eager result bitwise, and the reference within the band.""" + x, gate_up_weight, down_weight = _build(C1.spec, seed=_seed("graph")) + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(2): + op(x, gate_up_weight, down_weight) + torch.cuda.current_stream().wait_stream(side) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(x, gate_up_weight, down_weight) + + fresh, _, _ = _build(C1.spec, seed=_seed("graph replay")) + x.copy_(fresh) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(captured, op(x, gate_up_weight, down_weight)) + _compare(captured, op.reference(x, gate_up_weight, down_weight)) + + +@pytest.mark.parametrize( + "rows, hidden, inter, served", + [ + (1, HIDDEN, INTER_PER_RANK, True), + (8, HIDDEN, INTER_PER_RANK, True), + (0, HIDDEN, INTER_PER_RANK, False), + (9, HIDDEN, INTER_PER_RANK, False), + (8, HIDDEN, 2 * INTER_PER_RANK, False), + (8, HIDDEN, 4 * INTER_PER_RANK, False), + (8, HIDDEN, 16 * INTER_PER_RANK, False), + (8, 7168, INTER_PER_RANK, False), + ], + ids=["1 row", "8 rows", "0 rows", "9 rows", "tp8", "tp4", "tp1", "hidden 7168"], +) +def test_the_supported_predicate(rows, hidden, inter, served) -> None: + """The host-side answer the caller routes on: 1-8 rows of the TP16 slice + only. TP4 is what a dense MLP capped at the node's 4 GPUs would hand over.""" + assert torch.ops.trtllm.glm_dense_mlp_small_m_supported(rows, hidden, inter) is served + + +def _strided_x(): + x, w13, w2 = _build(C1.spec, seed=_seed("strided")) + wide = torch.zeros(x.shape[0], 2 * HIDDEN, dtype=x.dtype, device=x.device) + wide[:, :HIDDEN] = x + return wide[:, :HIDDEN], w13, w2 + + +def _misaligned_x(): + x, w13, w2 = _build(C1.spec, seed=_seed("misaligned")) + flat = torch.zeros(x.numel() + 1, dtype=x.dtype, device=x.device) + shifted = flat[1:].view(x.shape) + shifted.copy_(x) + return shifted, w13, w2 + + +def _too_many_rows(): + x, w13, w2 = _build(dict(C1.spec, rows=9), seed=_seed("9 rows")) + return x, w13, w2 + + +def _tp8_slice(): + x, _, _ = _build(C1.spec, seed=_seed("tp8")) + w13 = torch.zeros(4 * INTER_PER_RANK, HIDDEN, dtype=torch.bfloat16, device="cuda") + w2 = torch.zeros(HIDDEN, 2 * INTER_PER_RANK, dtype=torch.bfloat16, device="cuda") + return x, w13, w2 + + +def _fp16_x(): + x, w13, w2 = _build(C1.spec, seed=_seed("fp16")) + return x.to(torch.float16), w13, w2 + + +def _mismatched_down(): + x, w13, w2 = _build(C1.spec, seed=_seed("mismatch")) + return x, w13, w2[:, : INTER_PER_RANK // 2].contiguous() + + +@pytest.mark.parametrize( + "make", + [_strided_x, _misaligned_x, _too_many_rows, _tp8_slice, _fp16_x, _mismatched_down], + ids=lambda f: f.__name__.lstrip("_"), +) +def test_inputs_outside_the_domain_raise(make) -> None: + """Loud, never a silently wrong answer: these are why `is_valid` has + nothing left to refuse.""" + x, gate_up_weight, down_weight = make() + with pytest.raises(RuntimeError): + with validating(op): + op(x, gate_up_weight, down_weight) From ff3f54a9eb2d389447952b0195e1d4d9dd158480 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 15:51:01 -0700 Subject: [PATCH 08/53] Add the GLM-5.2 one-kernel BF16 embedding lookup op torch.ops.trtllm.glm_fused_embedding(input, weight) looks up int32 or int64 token ids of any shape in an unsharded BF16 table [vocab, hidden] with one kernel (kernels/glmMoeDsa/glmFusedEmbeddingLookup.cu): one thread per 16-byte vector, a grid of tokens x hidden / 1024 CTAs of 128 threads. An id outside [0, vocab), such as the -1 that pads rejected draft tokens, gives a zero row, the guard of the Embedding module's unsharded path. Rows are bit copies of the table, so the output equals the module's masked F.embedding chain bit for bit. Under programmatic dependent launch the kernel waits for its producer before it reads the ids and releases its dependents right after the wait; it touches no memory before it. The GLM-5.2 TP16 target keeps the [154880, 6144] table whole on every rank, so the lookup needs no collective. Its entry is ops/embedding/fused_embedding.py, certified bitwise on sm_107 by its test. Signed-off-by: Dongfeng Yu --- .../glmMoeDsa/glmFusedEmbeddingLookup.cu | 95 ++++++++++ .../glmMoeDsa/glmFusedEmbeddingLookup.h | 46 +++++ cpp/tensorrt_llm/thop/CMakeLists.txt | 1 + cpp/tensorrt_llm/thop/glmFusedEmbeddingOp.cpp | 74 ++++++++ .../ops/embedding/fused_embedding.py | 102 +++++++++++ .../_torch/custom_ops/cpp_custom_ops.py | 4 + .../test_modeling_v2_fused_embedding.py | 167 ++++++++++++++++++ 7 files changed, 489 insertions(+) create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.cu create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.h create mode 100644 cpp/tensorrt_llm/thop/glmFusedEmbeddingOp.cpp create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/embedding/fused_embedding.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/embedding/test_modeling_v2_fused_embedding.py diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.cu new file mode 100644 index 000000000000..282fd4fe89a7 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.cu @@ -0,0 +1,95 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/envUtils.h" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ +namespace +{ + +constexpr int kThreadsPerBlock = 128; +// BF16 elements per 16-byte vector. +constexpr int kElemsPerVec = 8; + +// blockIdx.x selects the token, blockIdx.y the 128-vector slice of its row; each thread copies one vector. +template +__global__ void __launch_bounds__(kThreadsPerBlock) fusedEmbeddingLookupKernel(uint4* __restrict__ out, + IdT const* __restrict__ ids, uint4 const* __restrict__ table, int64_t vocabSize, int32_t vecsPerRow) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + // The ids may be written by the preceding kernel. Dependents read the output only after their own wait. + cudaGridDependencySynchronize(); + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + int32_t const vec = static_cast(blockIdx.y * kThreadsPerBlock + threadIdx.x); + if (vec >= vecsPerRow) + { + return; + } + int64_t const token = blockIdx.x; + int64_t const id = static_cast(ids[token]); + + // Ids outside [0, vocabSize), such as the -1 padding of rejected draft tokens, give a zero row. + uint4 value = make_uint4(0U, 0U, 0U, 0U); + if (id >= 0 && id < vocabSize) + { + value = table[id * vecsPerRow + vec]; + } + out[token * vecsPerRow + vec] = value; +} + +} // namespace + +void invokeFusedEmbeddingLookup(__nv_bfloat16* out, void const* ids, bool idsAreInt64, __nv_bfloat16 const* table, + int64_t numTokens, int64_t vocabSize, int64_t hiddenSize, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(hiddenSize % kElemsPerVec == 0, "fusedEmbeddingLookup needs a hidden size divisible by 8"); + if (numTokens == 0) + { + return; + } + + int32_t const vecsPerRow = static_cast(hiddenSize / kElemsPerVec); + int32_t const blocksPerRow = (vecsPerRow + kThreadsPerBlock - 1) / kThreadsPerBlock; + dim3 const grid(static_cast(numTokens), static_cast(blocksPerRow)); + dim3 const block(kThreadsPerBlock); + auto* outVec = reinterpret_cast(out); + auto const* tableVec = reinterpret_cast(table); + if (idsAreInt64) + { + tensorrt_llm::common::launchWithPdlWhenEnabled("fusedEmbeddingLookup", fusedEmbeddingLookupKernel, + grid, block, 0, stream, outVec, static_cast(ids), tableVec, vocabSize, vecsPerRow); + } + else + { + tensorrt_llm::common::launchWithPdlWhenEnabled("fusedEmbeddingLookup", fusedEmbeddingLookupKernel, + grid, block, 0, stream, outVec, static_cast(ids), tableVec, vocabSize, vecsPerRow); + } +} + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.h new file mode 100644 index 000000000000..d15db15454f9 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.h @@ -0,0 +1,46 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +//! Embedding lookup of an unsharded BF16 table with the out-of-range guard of the Embedding module: +//! out[t, :] = (0 <= ids[t] < vocabSize) ? table[ids[t], :] : 0 +//! +//! One thread copies one 16-byte vector of a row; the grid covers numTokens x (hiddenSize / 8) vectors. +//! Under PDL the kernel waits for its producer before reading the ids and then lets its dependents launch. +//! +//! \param out Output [numTokens, hiddenSize], row-major, 16-byte aligned. +//! \param ids Token ids [numTokens], int64 when idsAreInt64, else int32. +//! \param table Embedding table [vocabSize, hiddenSize], row-major, 16-byte aligned. +//! \param hiddenSize Row length in elements; must be a multiple of 8. +void invokeFusedEmbeddingLookup(__nv_bfloat16* out, void const* ids, bool idsAreInt64, __nv_bfloat16 const* table, + int64_t numTokens, int64_t vocabSize, int64_t hiddenSize, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index ce6ad1d3eae7..e56c5cb31586 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -109,6 +109,7 @@ add_th_common_family( mxfp8Gemm.cpp fusedCatFp4Op.cpp fusedCatFp8Op.cpp + glmFusedEmbeddingOp.cpp weightOnlyQuantGemm.cpp weightOnlyQuantOp.cpp finegrained_mixed_dtype_gemm_thop.cpp diff --git a/cpp/tensorrt_llm/thop/glmFusedEmbeddingOp.cpp b/cpp/tensorrt_llm/thop/glmFusedEmbeddingOp.cpp new file mode 100644 index 000000000000..cd95ab8210e7 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmFusedEmbeddingOp.cpp @@ -0,0 +1,74 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +// Lookup of an unsharded BF16 embedding table in one kernel. Ids outside [0, weight.size(0)) give a zero row, +// the same guard as tensorrt_llm._torch.modules.embedding.pre_comm_embedding_ops. +at::Tensor glmFusedEmbedding(at::Tensor const& input, at::Tensor const& weight) +{ + CHECK_TH_CUDA(input); + CHECK_TH_CUDA(weight); + TORCH_CHECK(input.device() == weight.device(), "glm_fused_embedding: input and weight must be on the same device"); + TORCH_CHECK(input.scalar_type() == at::ScalarType::Int || input.scalar_type() == at::ScalarType::Long, + "glm_fused_embedding expects int32 or int64 ids, got ", input.scalar_type()); + CHECK_TYPE(weight, at::ScalarType::BFloat16); + CHECK_CONTIGUOUS(weight); + TORCH_CHECK(weight.dim() == 2, "glm_fused_embedding expects a 2D weight, got ", weight.dim(), "D"); + + int64_t const vocabSize = weight.size(0); + int64_t const hiddenSize = weight.size(1); + TORCH_CHECK(hiddenSize % 8 == 0, "glm_fused_embedding needs a hidden size divisible by 8, got ", hiddenSize); + TORCH_CHECK( + reinterpret_cast(weight.data_ptr()) % 16 == 0, "glm_fused_embedding needs a 16-byte aligned weight"); + + at::Tensor const ids = input.contiguous(); + std::vector outShape(input.sizes().begin(), input.sizes().end()); + outShape.push_back(hiddenSize); + at::Tensor const out = at::empty(outShape, weight.options()); + + auto const stream = at::cuda::getCurrentCUDAStream(input.get_device()); + tensorrt_llm::kernels::glm_moe_dsa::invokeFusedEmbeddingLookup(static_cast<__nv_bfloat16*>(out.data_ptr()), + ids.data_ptr(), ids.scalar_type() == at::ScalarType::Long, static_cast<__nv_bfloat16 const*>(weight.data_ptr()), + ids.numel(), vocabSize, hiddenSize, stream); + return out; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("glm_fused_embedding(Tensor input, Tensor weight) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_fused_embedding", &tensorrt_llm::torch_ext::glmFusedEmbedding); +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/embedding/fused_embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/embedding/fused_embedding.py new file mode 100644 index 000000000000..746926a8b680 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/embedding/fused_embedding.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Token embedding lookup from a replicated BF16 table in one kernel, via the trtllm +glm_fused_embedding op.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: GLM-5.2's vocabulary and hidden size: the table every rank holds whole. +VOCAB = 154880 +HIDDEN = 6144 + + +class GlmFusedEmbedding(OpWrapper): + """`out[..., :] = weight[input_ids[...], :]`, and a zero row for an id outside + `[0, weight.shape[0])`. + + The guard is the Embedding module's for an unsharded table: the -1 that + pads rejected draft tokens, or any other out-of-range id, reads as a zero + row instead of out of bounds. Rows are bit copies of the table. Returns a + new tensor of shape `input_ids.shape + (hidden,)` in the table's dtype; no + input is modified. + + `weight` is the whole table, [vocab, hidden], the same on every rank: the + lookup is local and needs no collective. A vocab-parallel shard is not a + valid `weight` (see `note`). + + One kernel on the current stream: one thread per 16-byte vector, a grid of + tokens x hidden / 1024 CTAs of 128 threads (8 x 6 at the c1 verify + forward). Under programmatic dependent launch it waits for its producer + before reading the ids, then releases its dependents at once and writes + `out`, so the kernel after it must wait before reading `out`. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The target looks up the verify forward's ids (one request's last accepted + # token and its 7 drafts at c1) in the replicated table; the engine's ids + # are int32. Prefill forwards (gen-server warm-up, PrefillTarget) look up a + # whole chunk. + CELLS: tuple[Cell, ...] = ( + Cell( + why="c1 verify forward: 8 int32 ids into the replicated 154880 x 6144 table", + spec=dict(tokens=8, ids_dtype=torch.int32, ids="random"), + ), + Cell( + why="decode without speculation: 1 id, a grid of 1 x 6 CTAs", + spec=dict(tokens=1, ids_dtype=torch.int32, ids="random"), + ), + Cell( + why="rejected-draft padding: -1 and ids >= vocab read as zero rows; rows 0, vocab-1", + spec=dict(tokens=8, ids_dtype=torch.int32, ids="edges"), + ), + Cell( + why="int64 ids: the second id type the kernel is instantiated for", + spec=dict(tokens=8, ids_dtype=torch.int64, ids="edges"), + ), + Cell( + why="prefill chunk: 8192 ids, a grid of 8192 x 6 CTAs", + spec=dict(tokens=8192, ids_dtype=torch.int32, ids="random"), + ), + ) + + note = """ + `weight`: BF16, contiguous, 2-D, 16-byte aligned, hidden a multiple of 8; + ids int32 or int64 on the table's device, any shape (made contiguous). The + op raises on anything else. + + A vocab-parallel shard passed as `weight` is not detectable from the + tensors and answers wrongly: the op indexes it with the global ids (rows + of the wrong tokens, zeros past the shard) instead of masking to the + shard's range and summing over ranks. The target keeps the table whole. + + Bitwise equal to the Embedding module's masked F.embedding chain for an + unsharded table. + """ + + def __call__(self, input_ids: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + return torch.ops.trtllm.glm_fused_embedding(input_ids, weight) + + def reference(self, input_ids: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """The table rows of the in-range ids, zero rows elsewhere: a gather, + nothing accumulated or rounded.""" + valid = (input_ids >= 0) & (input_ids < weight.shape[0]) + rows = torch.where(valid, input_ids, torch.zeros_like(input_ids)).long() + return torch.nn.functional.embedding(rows, weight).masked_fill(~valid.unsqueeze(-1), 0) + + def is_valid(self, input_ids: torch.Tensor, weight: torch.Tensor) -> None: + """Nothing to refuse: every input the op accepts it computes as stated. + + The op raises on every dtype, layout, alignment and device it does not + serve. What is left is whether `weight` is the whole table (see + `note`), which no check on the tensors can see. + """ + return None + + +fused_embedding = GlmFusedEmbedding() diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index fd8961fa4711..6a998e291ae7 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -985,6 +985,10 @@ def _(pe: torch.Tensor, nope: torch.Tensor): scale = pe.new_empty((M, 1), dtype=torch.int32) return packed, scale + @torch.library.register_fake("trtllm::glm_fused_embedding") + def _(input: torch.Tensor, weight: torch.Tensor): + return weight.new_empty((*input.shape, weight.shape[1])) + @torch.library.register_fake("trtllm::glm_dense_mlp_small_m") def _(x, gate_up_weight, down_weight, gate_up_tiling=-1, down_tiling=-1): return x.new_empty((x.shape[0], down_weight.shape[0])) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/embedding/test_modeling_v2_fused_embedding.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/embedding/test_modeling_v2_fused_embedding.py new file mode 100644 index 000000000000..555e0509d8dc --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/embedding/test_modeling_v2_fused_embedding.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM-5.2 tp16 fused_embedding entry (sm_107). + +A gather, so everything is gated bitwise: the cells against the entry's +reference, and the op against the Embedding module's own masked F.embedding +chain, which it replaces. Then CUDA-graph replay, id shapes, and the inputs the +op refuses. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.embedding.fused_embedding import ( # noqa: E501 + HIDDEN, + VOCAB, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.embedding.fused_embedding import ( # noqa: E501 + fused_embedding as op, +) +from tensorrt_llm._torch.modules.embedding import pre_comm_embedding_ops +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +assert torch.cuda.is_available(), "fused_embedding requires a CUDA device" + +#: model.embed_tokens.weight's standard deviation in the checkpoint. +TABLE_STD = 0.0076 + + +def _seed(text: str) -> int: + return zlib.crc32(text.encode()) % 2**31 + + +@pytest.fixture(scope="module") +def table(): + """The replicated [154880, 6144] bf16 table (1.9 GB), built once per module.""" + g = torch.Generator(device="cuda").manual_seed(_seed("table")) + weight = torch.randn(VOCAB, HIDDEN, generator=g, device="cuda", dtype=torch.bfloat16) + yield weight.mul_(TABLE_STD) + del weight + torch.cuda.empty_cache() + + +def _ids(spec, seed): + """`random`: uniform in-vocab ids. `edges`: the first and last rows, the -1 + padding, ids at and past the vocab, the id type's extremes, then random.""" + g = torch.Generator(device="cuda").manual_seed(seed) + ids = torch.randint(0, VOCAB, (spec["tokens"],), generator=g, device="cuda") + if spec["ids"] == "edges": + if spec["ids_dtype"] == torch.int64: + extremes = [2**33, -(2**40)] + else: + extremes = [2**31 - 1, -(2**31)] + edges = torch.tensor([-1, 0, VOCAB - 1, VOCAB, VOCAB + 1000, *extremes], device="cuda") + ids[: edges.numel()] = edges + return ids.to(spec["ids_dtype"]) + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Bitwise, compared as raw 16-bit patterns: rows are copies, so there is + no tolerance to give, and a raw compare also holds for -0.0 and NaN rows.""" + assert out.dtype == ref.dtype and out.shape == ref.shape, (out.shape, ref.shape) + assert torch.equal(out.view(torch.int16), ref.view(torch.int16)), "rows differ" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying, table) -> None: # noqa: F811 — fixture by name + certifying(op) + ids = _ids(cell.spec, seed=_seed(cell.why)) + before = ids.clone() + out = op(ids, table) + assert out.shape == (*ids.shape, HIDDEN) and out.dtype == table.dtype + assert out.is_contiguous() + assert torch.equal(ids, before), "the ids are read-only" + _compare(out, op.reference(ids, table)) + + +@pytest.mark.parametrize("ids_dtype", [torch.int32, torch.int64], ids=["int32", "int64"]) +def test_bitwise_equal_to_the_embedding_module_chain(ids_dtype, table) -> None: + """The masked F.embedding chain the Embedding module runs for an unsharded + table, with its out-of-range guard: the op computes exactly that.""" + ids = _ids(dict(tokens=64, ids_dtype=ids_dtype, ids="edges"), seed=_seed("chain")) + chain = pre_comm_embedding_ops( + ids, + table, + tp_size=1, + tp_rank=0, + tp_mode=None, + vocab_start_index=0, + vocab_end_index=VOCAB, + gather_output=False, + padding_size=0, + ) + _compare(op(ids, table), chain) + + +def test_cuda_graph_capture_and_replay(table) -> None: + """The verify forward's lookup runs inside the step's CUDA graph: replay + reads the ids the buffer holds at replay time.""" + spec = op.CELLS[0].spec + ids = _ids(spec, seed=_seed("graph")) + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(2): + op(ids, table) + torch.cuda.current_stream().wait_stream(side) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(ids, table) + + for seed, kind in (("graph random", "random"), ("graph edges", "edges")): + ids.copy_(_ids(dict(spec, ids=kind), seed=_seed(seed))) + graph.replay() + torch.cuda.synchronize() + _compare(captured, op.reference(ids, table)) + + +def test_the_output_follows_the_id_shape(table) -> None: + """Any id shape gives `ids.shape + (hidden,)`; strided ids are made + contiguous first; no ids give an empty result without a launch.""" + ids = _ids(dict(tokens=12, ids_dtype=torch.int64, ids="edges"), seed=_seed("shape")) + flat = op(ids, table) + grid = op(ids.view(3, 4), table) + assert grid.shape == (3, 4, HIDDEN) + _compare(grid.view(12, HIDDEN), flat) + _compare(op(ids[::2], table), flat[::2].contiguous()) + assert op(ids[:0], table).shape == (0, HIDDEN) + + +def _small_table(dtype=torch.bfloat16, hidden=64): + return torch.randn(32, hidden, device="cuda").to(dtype) + + +def _small_ids(): + return torch.arange(4, device="cuda", dtype=torch.int32) + + +@pytest.mark.parametrize( + "make", + [ + lambda: (_small_ids().float(), _small_table()), + lambda: (_small_ids().cpu(), _small_table()), + lambda: (_small_ids(), _small_table(torch.float16)), + lambda: (_small_ids(), _small_table(hidden=128)[:, :64]), + lambda: (_small_ids(), _small_table(hidden=60)), + ], + ids=["float ids", "cpu ids", "fp16 table", "strided table", "hidden 60"], +) +def test_inputs_outside_the_domain_raise(make) -> None: + """Loud, never a silently wrong answer: these are why `is_valid` has + nothing left to refuse.""" + input_ids, weight = make() + with pytest.raises(RuntimeError): + with validating(op): + op(input_ids, weight) From 4cac427a1ddceccc7700cd8c8c7a3c47ef25fe43 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 15:57:42 -0700 Subject: [PATCH 09/53] Run the GLM-5.2 decode step's embedding and dense MLPs on the GLM ops The verify forward looks its ids up in the whole BF16 table with one kernel, and the dense layers 0-2 run their 768 intermediate columns per rank as two weight-streaming kernels for up to 8 rows; the post-MLP all-reduce sums the rank partials as before. A batch the dense kernels do not serve takes the generic gated MLP. Signed-off-by: Dongfeng Yu --- .../glm_5_2_nvfp4__sm_107__tp16/modeling.py | 28 +++++++++++++++---- 1 file changed, 22 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py index 38d3e1d3d766..4b356ab23023 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py @@ -28,6 +28,8 @@ from tensorrt_llm._torch.speculative import SpecMetadata from . import weights as _weights +from .ops.embedding.fused_embedding import fused_embedding +from .ops.gemm.dense_mlp_small_m import INTER_PER_RANK, dense_mlp_small_m # The deployment this target is written for; route() selects it on the same # facts, and the constructor states them again for a direct construction. @@ -76,6 +78,11 @@ class DecodeTarget: - Layers 3-77 (MoE): both all-reduces carry the residual add and the following norm; the experts finalize before theirs (16 ranks span nodes). - The last all-reduce applies the model's final norm (``next_layer_layernorm``). + + The embedding table is whole on every rank, so the lookup is one local + kernel. The dense MLP runs as two weight-streaming kernels on this rank's + 768 intermediate columns for up to 8 rows (every c1 step: one request and + its 7 drafts); a larger batch takes the generic gated MLP. """ def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: @@ -110,13 +117,16 @@ def forward( ) -> torch.Tensor: core = self.core if inputs_embeds is None: - inputs_embeds = core.embed_tokens(input_ids) + inputs_embeds = fused_embedding(input_ids, core.embed_tokens.weight) residual = inputs_embeds hidden_states = core.layers[0].input_layernorm(inputs_embeds) + small_m = torch.ops.trtllm.glm_dense_mlp_small_m_supported( + hidden_states.shape[0], core.config.hidden_size, INTER_PER_RANK + ) for layer in core.layers[: core.num_hidden_layers]: if isinstance(layer.mlp, GatedMLP): hidden_states, residual = self._dense_layer( - layer, position_ids, hidden_states, residual, attn_metadata + layer, position_ids, hidden_states, residual, attn_metadata, small_m ) else: hidden_states, residual = self._moe_layer( @@ -125,7 +135,7 @@ def forward( return hidden_states @staticmethod - def _dense_layer(layer, position_ids, hidden_states, residual, attn_metadata): + def _dense_layer(layer, position_ids, hidden_states, residual, attn_metadata, small_m): hidden_states = layer.self_attn( position_ids=position_ids, hidden_states=hidden_states, @@ -133,9 +143,15 @@ def _dense_layer(layer, position_ids, hidden_states, residual, attn_metadata): all_reduce_params=AllReduceParams(enable_allreduce=True), ) hidden_states, residual = layer.post_attention_layernorm(hidden_states, residual) - hidden_states = layer.mlp( - hidden_states, final_all_reduce_params=AllReduceParams(enable_allreduce=False) - ) + if small_m: + # The rank partial; the all-reduce below sums it over the 16 ranks. + hidden_states = dense_mlp_small_m( + hidden_states, layer.mlp.gate_up_proj.weight, layer.mlp.down_proj.weight + ) + else: + hidden_states = layer.mlp( + hidden_states, final_all_reduce_params=AllReduceParams(enable_allreduce=False) + ) return layer.allreduce( hidden_states, all_reduce_params=AllReduceParams( From 712dddbb9e4b53e29b7f9cedc12ee6f94d9ec79b Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 15:53:58 -0700 Subject: [PATCH 10/53] Add the GLM-5 fused-A GEMM as the glm_fused_a_gemm op kernels/glmMoeDsa/glmFusedAGemm.cu computes kv_a_proj_with_mqa for GLM-5 at hidden 6144 and up to 8 tokens: 2624 output rows (q_a, compressed KV, k_pe), or 2784 with the DSA indexer's wk and weights_proj rows appended to the same weight. One CTA per 16 weight rows (164 or 174 CTAs, one per SM), and all eight warps issue the HMMAs: warp w owns the 256-wide k stages w, w + 8 and w + 16. Before the grid-dependency wait each CTA stages its weights by tensor copies: stage w into shared memory, stages w + 8 and w + 16 through shared memory into registers, warp w issuing its register-stage copies w * 250 ns after its first copy. Then it prefetches up to four caller tensors (weights that later kernels of the step read) into L2 in 16 KB bulk prefetches spread over the CTAs. After the wait it copies the activation tiles and triggers its dependents. The FP32 sums of the eight warps are added in warp order and rounded to BF16 once. torch.ops.trtllm.glm_fused_a_gemm(hidden_states, weight, l2_prefetch) takes the weight as stored and raises on shapes the kernel does not run. The modeling_v2 entry ops/gemm/fused_a_gemm.py of the glm_5_2_nvfp4 sm_107 tp16 target states the shape gates in is_valid; its sm_107 test drives the c1 cells and checks the 2624 columns of the 2784-row call, agreement with dsv3_fused_a_gemm_op, CUDA graph replay and the op's rejections. Signed-off-by: Dongfeng Yu --- .../kernels/glmMoeDsa/glmFusedAGemm.cu | 467 ++++++++++++++++++ .../kernels/glmMoeDsa/glmFusedAGemm.h | 64 +++ cpp/tensorrt_llm/thop/CMakeLists.txt | 1 + cpp/tensorrt_llm/thop/glmFusedAGemmOp.cpp | 130 +++++ .../ops/gemm/fused_a_gemm.py | 158 ++++++ .../gemm/test_modeling_v2_fused_a_gemm.py | 239 +++++++++ 6 files changed, 1059 insertions(+) create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.cu create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.h create mode 100644 cpp/tensorrt_llm/thop/glmFusedAGemmOp.cpp create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/fused_a_gemm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_fused_a_gemm.py diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.cu new file mode 100644 index 000000000000..3f8d4b4a959f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.cu @@ -0,0 +1,467 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaDriverWrapper.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +extern "C" +{ + __device__ uint32_t __nvvm_get_smem_pointer(void*); +} + +namespace +{ + +using bf16_t = __nv_bfloat16; + +// Warp w issues the tensor copies of its two register stages w * kFusedAMma8PaceNs after the copy of its shared-memory +// stage (1.75 us over the CTA's eight warps); the CTA's share of the L2 prefetch ranges follows once its staged tiles +// are in registers. +constexpr int kFusedAMma8PaceNs = 250; + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 +__device__ void initialize_barrier(uint64_t* smem_barrier, // 64 bits user-manged barrier in smem + int thread_count = 1) // Thread count expected to arrive/wait on this barrier +{ + uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier); + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(smem_int_ptr), "r"(thread_count)); +} + +// Barrier wait +__device__ void wait_barrier(uint64_t* smem_barrier, // 64 bits user-manged barrier in smem + int phase_bit) // Current phase bit the barrier waiting to flip +{ + uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier); + asm volatile( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra DONE;\n" + "bra LAB_WAIT;\n" + "DONE:\n" + "}\n" ::"r"(smem_int_ptr), + "r"(phase_bit)); +} + +// Prefetches the ranges into L2 in 16 KB pieces. Piece q of the concatenated ranges goes to CTA q % gridDim.x: an SM +// issues bulk prefetches at a limited rate, so every CTA takes an equal share. +__device__ void prefetch_ranges_l2(FusedAL2Prefetch const& ranges, int local_tid, int thread_cnt) +{ + constexpr size_t piece_bytes = 16384; + size_t piece_cnt[FusedAL2Prefetch::kMaxRanges]; + size_t total_piece_cnt = 0; +#pragma unroll + for (int i = 0; i < FusedAL2Prefetch::kMaxRanges; i++) + { + piece_cnt[i] = (ranges.bytes[i] + piece_bytes - 1) / piece_bytes; + total_piece_cnt += piece_cnt[i]; + } + for (size_t q = blockIdx.x + static_cast(local_tid) * gridDim.x; q < total_piece_cnt; + q += static_cast(thread_cnt) * gridDim.x) + { + size_t piece_idx = q; + int range_idx = 0; + while (piece_idx >= piece_cnt[range_idx]) + { + piece_idx -= piece_cnt[range_idx]; + range_idx++; + } + size_t const offset = piece_idx * piece_bytes; + size_t const remaining = ranges.bytes[range_idx] - offset; + uint32_t const bytes = static_cast(remaining < piece_bytes ? remaining : piece_bytes) & ~15u; + if (bytes > 0) + { + char const* gmem_ptr = static_cast(ranges.ptr[range_idx]) + offset; + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;\n" ::"l"(gmem_ptr), "r"(bytes) : "memory"); + } + } +} + +// Tensor copies completing on an mbarrier, ldmatrix and HMMA on shared addresses. +__device__ __forceinline__ void mma8_expect_tx(uint64_t* smem_barrier, uint32_t bytes) +{ + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"( + __nvvm_get_smem_pointer(smem_barrier)), + "r"(bytes) + : "memory"); +} + +__device__ __forceinline__ void mma8_tma_load_2d( + void* smem_dst, CUtensorMap const* tensor_map, int c0, int c1, uint64_t* smem_barrier) +{ + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1, {%2, %3}], [%4];\n" :: + "r"(__nvvm_get_smem_pointer(smem_dst)), + "l"(reinterpret_cast(tensor_map)), "r"(c0), "r"(c1), "r"(__nvvm_get_smem_pointer(smem_barrier)) + : "memory"); +} + +__device__ __forceinline__ void mma8_tma_load_3d( + void* smem_dst, CUtensorMap const* tensor_map, int c0, int c1, int c2, uint64_t* smem_barrier) +{ + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1, {%2, %3, %4}], " + "[%5];\n" ::"r"(__nvvm_get_smem_pointer(smem_dst)), + "l"(reinterpret_cast(tensor_map)), "r"(c0), "r"(c1), "r"(c2), + "r"(__nvvm_get_smem_pointer(smem_barrier)) + : "memory"); +} + +__device__ __forceinline__ void mma8_ldsm_x4(uint32_t smem_addr, uint32_t (&r)[4]) +{ + asm volatile("ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(smem_addr)); +} + +__device__ __forceinline__ void mma8_hmma(float (&d)[4], uint32_t const (&a)[4], uint32_t b0, uint32_t b1) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +} + +__device__ __forceinline__ uint64_t mma8_globaltimer() +{ + uint64_t t; + asm volatile("mov.u64 %0, %%globaltimer;\n" : "=l"(t)); + return t; +} +#endif + +// The GLM-5 fused-A GEMM: gemm_k 6144, 16 weight rows per CTA, up to 8 tokens, with all eight warps issuing the +// HMMAs. output: [gemm_n, gemm_m]; tma_a: the tensor map of the weights [gemm_m, gemm_k]; tma_b: that of the +// activations [gemm_n, gemm_k] (makeFusedAMma8Map). +// +// Stage s is the k range [256 s, 256 s + 256); warp w owns stages w, w + 8 and w + 16 and waits only on its own data. +// Before the grid dependency wait (no kernel of the step writes the weights): the weight tile of stage w is copied into +// shared memory (one 3D tensor copy: four 64-wide k boxes of the CTA's 16 rows, 128-byte swizzle), the weights of +// stages w + 8 and w + 16 reach registers, and the l2_prefetch ranges go to L2. After the wait: the activation tile of +// each of the warp's stages (one 3D copy at 8 tokens, else four 2D boxes of gemm_n rows; tile rows past gemm_n only +// feed HMMA columns that are never stored). The dependent launch is triggered right after the wait. +// +// The register stages permute the k order inside each 32-wide k pair alike for both operands: lane (g = lane / 4, +// c = lane % 4) holds k [8 c, 8 c + 8) of weight rows g and g + 8 and reads k [8 c, 8 c + 8) of token g from the +// activation tile; HMMA step t of the pair takes words 2 t and 2 t + 1 of each. Each warp keeps two accumulators +// (16-wide k step i of a stage adds into i % 2); the eight warp partials are summed in warp order. The products are +// those of a plain GEMM, summed in this order in FP32 and rounded to BF16 once. +// +// The weights of stages w + 8 and w + 16 come in through shared memory, still before the wait: two more 3D tensor +// copies per warp into the space of the activation tiles (those land only after the wait) extended by 32 KB, read into +// registers once they land. Tensor copies leave the SM's load path to a co-resident kernel (the post-MoE all-reduce hop +// that fused-A follows). Lane 0 of warp w issues its two register-stage copies w * kFusedAMma8PaceNs after its stage-w +// copy (%globaltimer), and the l2_prefetch ranges go out only once the CTA's staged tiles are in registers: the CTAs' +// weights reach HBM over 7 * kFusedAMma8PaceNs ahead of the ranges, and the hop beside fused-A keeps more of the HBM +// bandwidth. Issue times do not change the outputs. +template +__global__ __launch_bounds__(256, 1) void fused_a_gemm_mma8_kernel(bf16_t* output, int gemm_n, + FusedAL2Prefetch l2_prefetch_ranges, const __grid_constant__ CUtensorMap tma_a, + const __grid_constant__ CUtensorMap tma_b) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + constexpr int gemm_k = kFusedAHiddenSize; + constexpr int warp_cnt = 8; + constexpr int tile_m = kFusedARowsPerCta; + constexpr int tile_n = kFusedAMaxTokens; + constexpr int tile_k = 256; + constexpr int box_k = 64; + constexpr int boxes = tile_k / box_k; + constexpr int pair_cnt = tile_k / 32; + constexpr int acc_cnt = 2; + static_assert(gemm_k == 3 * warp_cnt * tile_k, "warp w owns stages w, w + 8 and w + 16"); + static_assert(gemm_m % tile_m == 0); + constexpr uint32_t a_box_bytes = tile_m * box_k * sizeof(bf16_t); + constexpr uint32_t b_box_bytes = tile_n * box_k * sizeof(bf16_t); + constexpr uint32_t a_tile_bytes = boxes * a_box_bytes; + constexpr uint32_t b_tile_bytes = boxes * b_box_bytes; + + // 1024-byte aligned (the 128-byte swizzle atoms): stage barriers, 8 weight tiles, 24 activation tiles. Before the + // wait, the 16 register-stage weight tiles sit where the activation tiles land after it (warp w's stages w + 8 and + // w + 16 at tiles 2 w and 2 w + 1), barrier 24 + w for warp w's two. + extern __shared__ char smem[]; + char* smem_base = smem + ((1024u - (__nvvm_get_smem_pointer(smem) & 1023u)) & 1023u); + uint64_t* full = reinterpret_cast(smem_base); + uint64_t* staged = full + 3 * warp_cnt; + char* smem_a = smem_base + 1024; + char* smem_b = smem_a + warp_cnt * a_tile_bytes; + + int const warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + int const lane = threadIdx.x % 32; + int const m0 = blockIdx.x * tile_m; + int const g = lane / 4; + int const c = lane % 4; + uint32_t const b_tx_bytes = boxes * gemm_n * box_k * sizeof(bf16_t); + bool const b_3d = gemm_n == tile_n; + + if (threadIdx.x == 0) + { + for (int i = 0; i < 3 * warp_cnt; i++) + { + // Stage s's barrier: the weight and the activation copies (stages 0-7), the activation copy (8-23). + initialize_barrier(full + i, i < warp_cnt ? 2 : 1); + } + for (int i = 0; i < warp_cnt; i++) + { + initialize_barrier(staged + i, 1); + } + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + } + __syncthreads(); + + if (lane == 0) + { + asm volatile("prefetch.tensormap [%0];\n" ::"l"(reinterpret_cast(&tma_a)) : "memory"); + asm volatile("prefetch.tensormap [%0];\n" ::"l"(reinterpret_cast(&tma_b)) : "memory"); + mma8_expect_tx(full + warp, a_tile_bytes); + mma8_tma_load_3d(smem_a + warp * a_tile_bytes, &tma_a, 0, m0, warp * boxes, full + warp); + mma8_expect_tx(staged + warp, 2 * a_tile_bytes); + uint64_t const due = mma8_globaltimer() + static_cast(warp) * kFusedAMma8PaceNs; + while (mma8_globaltimer() < due) + { + } +#pragma unroll + for (int j = 0; j < 2; j++) + { + mma8_tma_load_3d(smem_b + (2 * warp + j) * a_tile_bytes, &tma_a, 0, m0, + (warp + warp_cnt * (j + 1)) * boxes, staged + warp); + } + } + uint4 a_reg[2][pair_cnt][2]; + // Every weight copy of the CTA is queued ahead of its share of the L2 prefetch list. + __syncthreads(); + // The staged tiles into registers: lane (g, c) takes k [8 c, 8 c + 8) + 32 p of rows g and g + 8, i.e. 16-byte + // chunk (p % 2) * 4 + c of box p / 2 (128-byte rows, 128-byte swizzle). + wait_barrier(staged + warp, 0); +#pragma unroll + for (int j = 0; j < 2; j++) + { + uint32_t const tile = __nvvm_get_smem_pointer(smem_b + (2 * warp + j) * a_tile_bytes); +#pragma unroll + for (int p = 0; p < pair_cnt; p++) + { +#pragma unroll + for (int h = 0; h < 2; h++) + { + uint4& v = a_reg[j][p][h]; + asm volatile("ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "r"(tile + (p / 2) * a_box_bytes + (g + 8 * h) * 128 + ((((p % 2) * 4 + c) ^ g) << 4))); + } + } + } + // The activation copies after the wait overwrite the staged tiles. + __syncthreads(); + if constexpr (l2_prefetch) + { + prefetch_ranges_l2(l2_prefetch_ranges, threadIdx.x, blockDim.x); + } + cudaGridDependencySynchronize(); + cudaTriggerProgrammaticLaunchCompletion(); + + if (lane == 0) + { + // The generic-proxy reads of the staged tiles before the tensor copies that overwrite them. + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); +#pragma unroll + for (int j = 0; j < 3; j++) + { + int const stage = warp + warp_cnt * j; + mma8_expect_tx(full + stage, b_tx_bytes); + if (b_3d) + { + mma8_tma_load_3d(smem_b + stage * b_tile_bytes, &tma_b, 0, 0, stage * boxes, full + stage); + } + else + { +#pragma unroll + for (int q = 0; q < boxes; q++) + { + mma8_tma_load_2d(smem_b + stage * b_tile_bytes + q * b_box_bytes, &tma_b, + stage * tile_k + q * box_k, 0, full + stage); + } + } + } + } + + float acc[acc_cnt][4] = {}; + // Stage w from shared memory. ldmatrix lanes: weight row (lane % 16) at k (lane / 16) * 8 of a 16-wide step; + // token (lane % 8) at k (lane / 8) * 8 of a pair of steps. + { + wait_barrier(full + warp, 0); + int const a_row = lane % 16; + int const a_k = (lane / 16) * 8; + int const b_row = lane % 8; + int const b_k = (lane / 8) * 8; + uint32_t const a_base = __nvvm_get_smem_pointer(smem_a + warp * a_tile_bytes); + uint32_t const b_base = __nvvm_get_smem_pointer(smem_b + warp * b_tile_bytes); +#pragma unroll + for (int s = 0; s < tile_k / 16; s += 2) + { + int const q = s / 4; + int const kk = (s % 4) * 16; + uint32_t b[4]; + mma8_ldsm_x4(b_base + q * b_box_bytes + b_row * 128 + ((((kk + b_k) >> 3) ^ (b_row & 7)) << 4), b); +#pragma unroll + for (int t = 0; t < 2; t++) + { + uint32_t a[4]; + mma8_ldsm_x4( + a_base + q * a_box_bytes + a_row * 128 + ((((kk + 16 * t + a_k) >> 3) ^ (a_row & 7)) << 4), a); + mma8_hmma(acc[(s + t) % acc_cnt], a, b[2 * t], b[2 * t + 1]); + } + } + } + // Stages w + 8 and w + 16: weights from registers, activations from shared memory in the same permuted k order. +#pragma unroll + for (int j = 0; j < 2; j++) + { + int const stage = warp + warp_cnt * (j + 1); + wait_barrier(full + stage, 0); + uint32_t const b_base = __nvvm_get_smem_pointer(smem_b + stage * b_tile_bytes); +#pragma unroll + for (int p = 0; p < pair_cnt; p++) + { + int const chunk = (p % 2) * 4 + c; + uint4 bq; + asm volatile("ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(bq.x), "=r"(bq.y), "=r"(bq.z), "=r"(bq.w) + : "r"(b_base + (p / 2) * b_box_bytes + g * 128 + ((chunk ^ (g & 7)) << 4))); + uint32_t const a0[4] = {a_reg[j][p][0].x, a_reg[j][p][1].x, a_reg[j][p][0].y, a_reg[j][p][1].y}; + uint32_t const a1[4] = {a_reg[j][p][0].z, a_reg[j][p][1].z, a_reg[j][p][0].w, a_reg[j][p][1].w}; + mma8_hmma(acc[(2 * p) % acc_cnt], a0, bq.x, bq.y); + mma8_hmma(acc[(2 * p + 1) % acc_cnt], a1, bq.z, bq.w); + } + } + + // Warp partials into the warp's consumed stage w + 8 activation tile (no copy is pending there), then the sum + // over the warps in warp order. Output (row r, token n) is register (n % 2) + 2 (r / 8) of lane + // (r % 8) * 4 + n / 2. + __syncwarp(); + float* partial = reinterpret_cast(smem_b + (warp + warp_cnt) * b_tile_bytes); +#pragma unroll + for (int r = 0; r < 4; r++) + { + partial[lane * 4 + r] = acc[0][r] + acc[1][r]; + } + __syncthreads(); + if (threadIdx.x < tile_m * tile_n) + { + int const r = threadIdx.x % tile_m; + int const n = threadIdx.x / tile_m; + int const src = ((r % 8) * 4 + n / 2) * 4 + (n % 2) + 2 * (r / 8); + float v = 0.f; +#pragma unroll + for (int w = 0; w < warp_cnt; w++) + { + v += reinterpret_cast(smem_b + (w + warp_cnt) * b_tile_bytes)[src]; + } + if (n < gemm_n) + { + output[static_cast(n) * gemm_m + m0 + r] = __float2bfloat16_rn(v); + } + } +#endif +} + +// Tensor map of a row-major BF16 [rows, cols] matrix with the 128-byte swizzle: dims (k within a 64-wide box, row, box) +// and box (64, box_rows, 4) when box3d (one copy loads [4 boxes][box_rows][64]), else dims (k, row) and box +// (64, box_rows). +CUtensorMap makeFusedAMma8Map(void const* ptr, int rows, int cols, int box_rows, bool box3d) +{ + CUtensorMap map; + cuuint64_t const dims3[3] = {64, static_cast(rows), static_cast(cols / 64)}; + cuuint64_t const strides3[2] = {static_cast(cols) * sizeof(bf16_t), 64 * sizeof(bf16_t)}; + cuuint32_t const box3[3] = {64, static_cast(box_rows), 4}; + cuuint64_t const dims2[2] = {static_cast(cols), static_cast(rows)}; + cuuint64_t const strides2[1] = {static_cast(cols) * sizeof(bf16_t)}; + cuuint32_t const box2[2] = {64, static_cast(box_rows)}; + cuuint32_t const elem_strides[3] = {1, 1, 1}; + TLLM_CU_CHECK(tensorrt_llm::common::CUDADriverWrapper::getInstance()->cuTensorMapEncodeTiled(&map, + CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, box3d ? 3 : 2, const_cast(ptr), box3d ? dims3 : dims2, + box3d ? strides3 : strides2, box3d ? box3 : box2, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, + CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_L2_256B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return map; +} + +// weight: [gemm_m, 6144]; hiddenStates: [numTokens, 6144], numTokens in [1, 8]. +template +void launchFusedAMma8(bf16_t* output, bf16_t const* hiddenStates, bf16_t const* weight, int numTokens, + FusedAL2Prefetch const* l2Prefetch, cudaStream_t stream) +{ + // Alignment slack, stage barriers, 8 weight tiles of 16 x 256, and the space of the 24 activation tiles of 8 x 256 + // extended to the 16 register-stage weight tiles of 16 x 256. + constexpr int kSmemBytes = 1024 + 1024 + 8 * 16 * 256 * sizeof(bf16_t) + 16 * 16 * 256 * sizeof(bf16_t); + CUtensorMap const tmaA = makeFusedAMma8Map(weight, gemm_m, kFusedAHiddenSize, kFusedARowsPerCta, true); + CUtensorMap const tmaB + = makeFusedAMma8Map(hiddenStates, numTokens, kFusedAHiddenSize, numTokens, numTokens == kFusedAMaxTokens); + auto* kernel = l2Prefetch != nullptr ? &fused_a_gemm_mma8_kernel + : &fused_a_gemm_mma8_kernel; + TLLM_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes)); + cudaLaunchConfig_t config; + config.gridDim = dim3(gemm_m / kFusedARowsPerCta, 1, 1); + config.blockDim = dim3(256); + config.dynamicSmemBytes = kSmemBytes; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = attrs; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, output, numTokens, + l2Prefetch != nullptr ? *l2Prefetch : FusedAL2Prefetch{}, tmaA, tmaB)); +} + +} // namespace + +void invokeFusedAGemm(__nv_bfloat16* output, __nv_bfloat16 const* hiddenStates, __nv_bfloat16 const* weight, + int numTokens, int outFeatures, FusedAL2Prefetch const* l2Prefetch, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(numTokens >= 1 && numTokens <= kFusedAMaxTokens, + "The GLM fused-A GEMM takes 1 to %d tokens, got %d", kFusedAMaxTokens, numTokens); + switch (outFeatures) + { + case kFusedAOutFeatures: + launchFusedAMma8(output, hiddenStates, weight, numTokens, l2Prefetch, stream); + break; + case kFusedAIndexerOutFeatures: + launchFusedAMma8(output, hiddenStates, weight, numTokens, l2Prefetch, stream); + break; + default: + TLLM_THROW("The GLM fused-A GEMM serves %d or %d output features, got %d", kFusedAOutFeatures, + kFusedAIndexerOutFeatures, outFeatures); + } +} + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.h new file mode 100644 index 000000000000..b198da9087ba --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.h @@ -0,0 +1,64 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +//! The GEMM's K: the GLM-5 hidden size, 24 k stages of 256, three per MMA warp. +constexpr int kFusedAHiddenSize = 6144; +//! Rows of the GLM-5 kv_a_proj_with_mqa weight: q_a (2048), compressed KV (512) and k_pe (64). +constexpr int kFusedAOutFeatures = 2624; +//! Rows of the same weight with the DSA indexer's wk (128) and weights_proj (32) rows appended. +constexpr int kFusedAIndexerOutFeatures = 2784; +//! Weight rows per CTA: the grid is outFeatures / 16 CTAs (164 or 174), one per SM. +constexpr int kFusedARowsPerCta = 16; +//! Tokens per call at most: one activation tile holds 8 token rows. +constexpr int kFusedAMaxTokens = 8; + +//! Byte ranges that the GEMM prefetches into L2 before its grid-dependency wait: weights that later kernels of the +//! step read. Unused slots have zero bytes. +struct FusedAL2Prefetch +{ + static constexpr int kMaxRanges = 4; + void const* ptr[kMaxRanges]; + size_t bytes[kMaxRanges]; +}; + +//! output[t, r] = sum_k hiddenStates[t, k] * weight[r, k]: BF16 operands, FP32 accumulation, BF16 output, all three +//! dense row-major. hiddenStates is [numTokens, 6144] with numTokens in [1, 8]; weight is [outFeatures, 6144] with +//! outFeatures kFusedAOutFeatures or kFusedAIndexerOutFeatures; output is [numTokens, outFeatures]. +//! +//! Launched with the programmatic-dependent-launch attribute when TRTLLM_ENABLE_PDL is set. Before its grid-dependency +//! wait the kernel reads only the weight (tensor copies into shared memory and registers) and prefetches the +//! l2Prefetch ranges into L2 (none when l2Prefetch is null); it reads hiddenStates and writes output after the wait, +//! and triggers its dependents right after the wait. Nothing launched before it in the step may write the weight. +void invokeFusedAGemm(__nv_bfloat16* output, __nv_bfloat16 const* hiddenStates, __nv_bfloat16 const* weight, + int numTokens, int outFeatures, FusedAL2Prefetch const* l2Prefetch, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index e56c5cb31586..faaa106a4995 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -99,6 +99,7 @@ add_th_common_family( fp8Quantize.cpp dsv3FusedAGemmOp.cpp glmDenseMlpSmallMOp.cpp + glmFusedAGemmOp.cpp deepseekV4QNormOp.cpp fusedAddRMSNormQuant.cpp fusedActivationQuant.cpp diff --git a/cpp/tensorrt_llm/thop/glmFusedAGemmOp.cpp b/cpp/tensorrt_llm/thop/glmFusedAGemmOp.cpp new file mode 100644 index 000000000000..2a33b5fc1e1f --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmFusedAGemmOp.cpp @@ -0,0 +1,130 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.h" + +#include +#include +#include +#include + +namespace tk = tensorrt_llm::kernels; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +using tk::glm_moe_dsa::FusedAL2Prefetch; + +constexpr int64_t kTmaAlignmentBytes = 16; + +bool isAligned(torch::Tensor const& tensor) +{ + return reinterpret_cast(tensor.data_ptr()) % kTmaAlignmentBytes == 0; +} + +// One L2 prefetch range per tensor: its contiguous bytes on the GEMM's device. +FusedAL2Prefetch makeL2Prefetch(torch::TensorList tensors, torch::Device const& device) +{ + TORCH_CHECK(tensors.size() <= static_cast(FusedAL2Prefetch::kMaxRanges), + "glm_fused_a_gemm prefetches at most ", FusedAL2Prefetch::kMaxRanges, " tensors, got ", tensors.size()); + FusedAL2Prefetch ranges{}; + for (size_t i = 0; i < tensors.size(); ++i) + { + torch::Tensor const& tensor = tensors[i]; + TORCH_CHECK(tensor.device() == device, "l2_prefetch tensors must be on the GEMM's device"); + TORCH_CHECK(tensor.is_contiguous(), "l2_prefetch tensors must be contiguous"); + TORCH_CHECK(isAligned(tensor), "l2_prefetch tensors must be 16-byte aligned"); + ranges.ptr[i] = tensor.data_ptr(); + ranges.bytes[i] = static_cast(tensor.numel()) * tensor.element_size(); + } + return ranges; +} + +} // namespace + +//! The GLM-5 fused-A GEMM (kv_a_proj_with_mqa) for up to 8 tokens. +//! +//! @param hiddenStates BF16 CUDA tensor [num_tokens, 6144], contiguous, num_tokens in [1, 8]. +//! @param weight BF16 CUDA tensor [out_features, 6144], contiguous: the kv_a_proj_with_mqa weight (2624 rows: q_a, +//! compressed KV, k_pe) or that weight with the DSA indexer's wk and weights_proj rows appended (2784 rows). +//! @param l2Prefetch Up to four contiguous, 16-byte aligned tensors on the same device that the kernel prefetches into +//! L2 before its grid-dependency wait (weights later kernels of the step read); may be empty. +//! @return BF16 tensor [num_tokens, out_features] = hiddenStates @ weight^T, FP32 accumulation. +torch::Tensor glmFusedAGemm( + torch::Tensor const& hiddenStates, torch::Tensor const& weight, torch::TensorList l2Prefetch) +{ + namespace glm = tk::glm_moe_dsa; + TORCH_CHECK(hiddenStates.is_cuda() && weight.is_cuda(), "glm_fused_a_gemm expects CUDA tensors"); + TORCH_CHECK(hiddenStates.device() == weight.device(), "hidden_states and weight must be on the same device"); + TORCH_CHECK(hiddenStates.scalar_type() == torch::kBFloat16 && weight.scalar_type() == torch::kBFloat16, + "glm_fused_a_gemm expects BF16 hidden_states and weight, got ", hiddenStates.scalar_type(), " and ", + weight.scalar_type()); + TORCH_CHECK(hiddenStates.dim() == 2 && weight.dim() == 2, + "glm_fused_a_gemm expects hidden_states [num_tokens, 6144] and weight [out_features, 6144]"); + TORCH_CHECK(hiddenStates.size(1) == glm::kFusedAHiddenSize && weight.size(1) == glm::kFusedAHiddenSize, + "glm_fused_a_gemm serves hidden size ", glm::kFusedAHiddenSize, ", got hidden_states ", hiddenStates.sizes(), + " and weight ", weight.sizes()); + TORCH_CHECK(weight.size(0) == glm::kFusedAOutFeatures || weight.size(0) == glm::kFusedAIndexerOutFeatures, + "glm_fused_a_gemm serves ", glm::kFusedAOutFeatures, " or ", glm::kFusedAIndexerOutFeatures, + " output features, got ", weight.size(0)); + TORCH_CHECK(hiddenStates.size(0) >= 1 && hiddenStates.size(0) <= glm::kFusedAMaxTokens, + "glm_fused_a_gemm takes 1 to ", glm::kFusedAMaxTokens, " tokens, got ", hiddenStates.size(0)); + TORCH_CHECK(hiddenStates.is_contiguous() && weight.is_contiguous(), + "glm_fused_a_gemm expects contiguous row-major hidden_states and weight"); + TORCH_CHECK(isAligned(hiddenStates) && isAligned(weight), + "glm_fused_a_gemm expects 16-byte aligned hidden_states and weight (tensor copies)"); + + c10::cuda::CUDAGuard const deviceGuard{hiddenStates.device()}; + TORCH_CHECK(tensorrt_llm::common::getSMVersion() >= 90, "glm_fused_a_gemm needs SM90 or later"); + int const numTokens = static_cast(hiddenStates.size(0)); + int const outFeatures = static_cast(weight.size(0)); + torch::Tensor output = torch::empty({numTokens, outFeatures}, hiddenStates.options()); + auto const stream = at::cuda::getCurrentCUDAStream(hiddenStates.get_device()); + auto* outputPtr = reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()); + auto const* hiddenStatesPtr = reinterpret_cast<__nv_bfloat16 const*>(hiddenStates.data_ptr()); + auto const* weightPtr = reinterpret_cast<__nv_bfloat16 const*>(weight.data_ptr()); + if (l2Prefetch.empty()) + { + glm::invokeFusedAGemm(outputPtr, hiddenStatesPtr, weightPtr, numTokens, outFeatures, nullptr, stream); + } + else + { + FusedAL2Prefetch const ranges = makeL2Prefetch(l2Prefetch, hiddenStates.device()); + glm::invokeFusedAGemm(outputPtr, hiddenStatesPtr, weightPtr, numTokens, outFeatures, &ranges, stream); + } + return output; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("glm_fused_a_gemm(Tensor hidden_states, Tensor weight, Tensor[] l2_prefetch) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_fused_a_gemm", &tensorrt_llm::torch_ext::glmFusedAGemm); +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/fused_a_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/fused_a_gemm.py new file mode 100644 index 000000000000..5f02dd788c34 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/fused_a_gemm.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GLM-5 fused-A GEMM (kv_a_proj_with_mqa, optionally with the DSA indexer's wk | weights_proj +rows) for up to 8 tokens, with an L2 prefetch of later weights, via the trtllm glm_fused_a_gemm +op.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +HIDDEN_SIZE = 6144 +# kv_a_proj_with_mqa: q_a (2048) | compressed KV (512) | k_pe (64). +OUT_FEATURES = 2624 +# The same rows, then the indexer's wk (128) and weights_proj (32) rows. +INDEXER_OUT_FEATURES = 2784 +ROWS_PER_CTA = 16 +MAX_TOKENS = 8 + +# What a c1 decode step lists behind fused-A, at TP16 per-rank shapes. +_INDEXER_WQ_B = ((4096, 2048), torch.bfloat16) # indexer wq_b, q_a norm folded in (not TP-split) +_Q_B_ABSORBED = ((2304, 2048), torch.bfloat16) # q_b with k_b absorbed: 4 heads x (512 + 64) +_V_B = ((4, 256, 512), torch.bfloat16) # v_b_proj, 4 heads +_O_PROJ = ((6144, 1024), torch.bfloat16) # o_proj shard +_LEADER_TOPK = ((8, 2048), torch.int32) # the indexer layer's top-k that a shared layer reuses + + +class GlmFusedAGemm(OpWrapper): + """`hidden_states @ weight.T`: BF16 operands, FP32 accumulation, BF16 out, `[tokens, rows]`. + + `weight` is the kv_a_proj_with_mqa weight as stored, `[2624, 6144]` row-major (q_a, compressed + KV, k_pe), or one tensor with the DSA indexer's wk and weights_proj rows appended, + `[2784, 6144]`. Output columns then read `[0, 2048)` q_a, `[2048, 2560)` compressed KV, + `[2560, 2624)` k_pe, `[2624, 2752)` indexer K, `[2752, 2784)` indexer weights; columns + `[0, 2624)` are bitwise those of the 2624-row call. + + One CTA per 16 weight rows (164 or 174 CTAs, one per SM) with all eight warps issuing the + HMMAs. Before its grid-dependency wait the kernel stages the CTA's weight rows (tensor copies + into shared memory, two thirds of them on into registers) and then prefetches the + `l2_prefetch` tensors into L2; it reads `hidden_states` only after the wait and triggers its + dependents right after the wait. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The c1 decode step runs the op 85 times: 78 verify layers at 8 rows plus 7 MTP draft steps. + # 22 calls (21 indexer layers and draft step 0) take the 2784-row weight. + CELLS: tuple[Cell, ...] = ( + Cell( + why="indexer layers' verify step and MTP draft step 0: 2784 rows, 8 tokens (3D " + "activation copies), the list indexer wq_b, v_b, o_proj", + spec=dict(tokens=8, out_features=2784, l2_prefetch=(_INDEXER_WQ_B, _V_B, _O_PROJ)), + ), + Cell( + why="shared-top-k layers' verify step (57 of 78): 2624 rows, 8 tokens, a full list " + "of four ending in the indexer layer's int32 top-k", + spec=dict( + tokens=8, + out_features=2624, + l2_prefetch=(_Q_B_ABSORBED, _V_B, _O_PROJ, _LEADER_TOPK), + ), + ), + Cell( + why="MTP draft steps 1-6: 2624 rows at one token (2D activation boxes of one row)", + spec=dict(tokens=1, out_features=2624, l2_prefetch=(_Q_B_ABSORBED, _V_B, _O_PROJ)), + ), + Cell( + why="4 tokens (MTP with 3 draft tokens): 2D boxes between 1 and 8 rows, 2784 rows", + spec=dict(tokens=4, out_features=2784, l2_prefetch=(_INDEXER_WQ_B, _V_B, _O_PROJ)), + ), + Cell( + why="empty list: the instance without L2 prefetches, for more than 32 rows per step", + spec=dict(tokens=8, out_features=2624, l2_prefetch=()), + ), + ) + + note = """ + `l2_prefetch` is a cache hint: up to four contiguous, 16-byte aligned CUDA tensors on the + GEMM's device, issued in 16 KB bulk prefetches spread over the CTAs once each CTA's weight + rows are staged. It never enters the output, and the kernel writes none of it. + + PDL (`TRTLLM_ENABLE_PDL`): before its grid-dependency wait the kernel reads `weight`, so no + kernel launched earlier in the step may write it (load-time weights only). Its dependents + launch once it passes the wait and must wait on their own grid dependency before reading the + output or writing `hidden_states`. + + Preconditions, enforced loudly by the op: CUDA BF16 2-D `hidden_states` and `weight` on one + device, contiguous and 16-byte aligned; hidden 6144; 2624 or 2784 rows; 1 to 8 tokens; at most + four `l2_prefetch` tensors, each contiguous, aligned and on the device; SM90 or later. + + Numerics: the products of a plain GEMM, summed in FP32 in the kernel's order (two + accumulators per warp, the eight warp partials in warp order) and rounded to BF16 once; + deterministic. The shared `dsv3_fused_a_gemm_op` sums the same products in another order: + nearly all outputs agree bitwise, the rest by one BF16 rounding step. + """ + + def __call__( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + l2_prefetch: list[torch.Tensor], + ) -> torch.Tensor: + return torch.ops.trtllm.glm_fused_a_gemm(hidden_states, weight, l2_prefetch) + + def reference( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + l2_prefetch: list[torch.Tensor], + ) -> torch.Tensor: + """`hidden_states @ weight.T` in FP32 (TF32 off), rounded to BF16 once. + + `l2_prefetch` does not enter the result; it keeps its name to mirror `__call__`. + """ + with true_fp32_matmul(): + return (hidden_states.float() @ weight.float().t()).to(hidden_states.dtype) + + def is_valid( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + l2_prefetch: list[torch.Tensor], + ) -> None: + # The shape gates of the kernel, which a target checks to route anything else to the + # shared dsv3_fused_a_gemm_op. The op raises on all but the last of them as well. + # gemm_k 6144: 24 k stages of 256, three per MMA warp. + assert hidden_states.dim() == 2 and hidden_states.shape[1] == HIDDEN_SIZE, ( + f"hidden_states must be [tokens, {HIDDEN_SIZE}]" + ) + # Tiles of 16 weight rows, one CTA each; instances exist for the 2624-row weight and the + # 2784-row weight with the indexer rows. + assert weight.dim() == 2 and weight.shape[1] == HIDDEN_SIZE, ( + f"weight must be [rows, {HIDDEN_SIZE}]" + ) + assert weight.shape[0] in (OUT_FEATURES, INDEXER_OUT_FEATURES), ( + f"weight must have {OUT_FEATURES} or {INDEXER_OUT_FEATURES} rows, got {weight.shape[0]}" + ) + # One activation tile of 8 token rows per k stage. + assert 1 <= hidden_states.shape[0] <= MAX_TOKENS, ( + f"the kernel takes 1 to {MAX_TOKENS} tokens, got {hidden_states.shape[0]}" + ) + # One CTA per SM in a single wave: the kernel stages its weights before the + # grid-dependency wait, and a second wave of CTAs would start only after the first ends. + # The output stays right; on fewer SMs the caller takes the shared dsv3_fused_a_gemm_op + # instead (32-row tiles, half the CTAs, for the 2624-row weight). + sms = torch.cuda.get_device_properties(hidden_states.device).multi_processor_count + ctas = weight.shape[0] // ROWS_PER_CTA + assert sms >= ctas, f"{ctas} CTAs need as many SMs for one wave, the device has {sms} SMs" + + +fused_a_gemm = GlmFusedAGemm() diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_fused_a_gemm.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_fused_a_gemm.py new file mode 100644 index 000000000000..eed24372acbe --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_fused_a_gemm.py @@ -0,0 +1,239 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM fused_a_gemm entry. + +The entry owns the contract. This file turns a cell into tensors and holds the claims no cell +carries: the 2624 columns of the 2784-row call, determinism and the list as a pure hint, +agreement with the shared fused-A kernel, CUDA graph replay, the op's own rejections, and the +guards. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.fused_a_gemm import ( # noqa: E501 + HIDDEN_SIZE, + INDEXER_OUT_FEATURES, + OUT_FEATURES, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.fused_a_gemm import ( # noqa: E501 + fused_a_gemm as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402,F401 — needs the path declared above; fixture by name + assert_within_ulp, + certifying, + validating, +) + +assert torch.cuda.is_available(), "fused_a_gemm requires a CUDA device" + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM fused-A GEMM is certified on sm_107 only" +) + +DEV = "cuda" +BF16 = torch.bfloat16 + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """One GEMM's worth of accumulation-order difference, then one BF16 rounding on each side. + + Kernel and reference sum the same products (a BF16 x BF16 product is exact in FP32) in FP32 + and round once to BF16; only the summation order differs. With randn operands the outputs sit + near sqrt(6144) ~ 80, the two FP32 sums differ by about 1e-5 there (FP32 rounding over a + few dozen additions of partials up to that size), and the BF16 spacing there is 0.5. So an + output either rounds to the same BF16 value or lands one rounding step away, and a step is + at most 2 ulp (2^-8) of the row's largest magnitude. + + So: 4 ulp element-wise against the row scale (twice the worst rounding step), 2 ulp of + relative RMS (a step on 1 % of the outputs would read as about 0.15). A wrong k stage, row + or token is off by a fraction of the row scale and misses both by orders. + """ + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0) + + +def _tensor(shape, dtype, generator): + if dtype.is_floating_point: + return torch.randn(shape, generator=generator, device=DEV).to(dtype) + return torch.randint(0, 1 << 20, shape, generator=generator, device=DEV, dtype=dtype) + + +def _build(spec, seed): + """Materialize one cell: randn activations and weight, and the list tensors at their shapes.""" + g = torch.Generator(device=DEV).manual_seed(seed) + hidden_states = _tensor((spec["tokens"], HIDDEN_SIZE), BF16, g) + weight = _tensor((spec["out_features"], HIDDEN_SIZE), BF16, g) + l2_prefetch = [_tensor(shape, dtype, g) for shape, dtype in spec["l2_prefetch"]] + return hidden_states, weight, l2_prefetch + + +# ── the certified cells ─────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture + certifying(op) + hidden_states, weight, l2_prefetch = _build( + cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31 + ) + snapshots = [t.clone() for t in l2_prefetch] + out = op(hidden_states, weight, l2_prefetch) + assert out.shape == (hidden_states.shape[0], weight.shape[0]) + assert out.dtype == BF16 and out.is_contiguous() + _compare(out, op.reference(hidden_states, weight, l2_prefetch)) + # The list is a cache hint: read, never written. + for tensor, snapshot in zip(l2_prefetch, snapshots): + assert torch.equal(tensor, snapshot) + + +# ── the claims no cell carries ──────────────────────────────────────────── + + +@pytest.mark.parametrize("tokens", [1, 8]) +def test_indexer_rows_leave_the_kv_a_columns_bitwise_unchanged(tokens) -> None: + """A target keeps one 2784-row tensor and reads kv_a_proj_with_mqa as its first 2624 rows; + every 16-row tile is computed alike in both grids.""" + g = torch.Generator(device=DEV).manual_seed(tokens) + hidden_states = _tensor((tokens, HIDDEN_SIZE), BF16, g) + extended = _tensor((INDEXER_OUT_FEATURES, HIDDEN_SIZE), BF16, g) + kv_a = extended[:OUT_FEATURES] + assert kv_a.is_contiguous() and kv_a.data_ptr() == extended.data_ptr() + out_extended = op(hidden_states, extended, []) + out_kv_a = op(hidden_states, kv_a, []) + assert torch.equal(out_extended[:, :OUT_FEATURES], out_kv_a) + + +def test_deterministic_and_the_list_does_not_enter_the_result() -> None: + g = torch.Generator(device=DEV).manual_seed(11) + hidden_states = _tensor((8, HIDDEN_SIZE), BF16, g) + weight = _tensor((INDEXER_OUT_FEATURES, HIDDEN_SIZE), BF16, g) + l2_prefetch = [_tensor((4096, 2048), BF16, g), _tensor((6144, 1024), BF16, g)] + first = op(hidden_states, weight, l2_prefetch) + second = op(hidden_states, weight, l2_prefetch) + without_list = op(hidden_states, weight, []) + assert torch.equal(first, second) + assert torch.equal(first, without_list) + assert first.data_ptr() != second.data_ptr() + + +@pytest.mark.parametrize("tokens", [1, 8]) +def test_agrees_with_the_shared_fused_a_kernel(tokens) -> None: + """dsv3_fused_a_gemm_op (32-row tiles, 4 MMA warps) sums the same products in another order: + nearly every output is bitwise equal, the rest one BF16 rounding step apart.""" + g = torch.Generator(device=DEV).manual_seed(100 + tokens) + hidden_states = _tensor((tokens, HIDDEN_SIZE), BF16, g) + weight = _tensor((OUT_FEATURES, HIDDEN_SIZE), BF16, g) + out = op(hidden_states, weight, []) + shared = torch.ops.trtllm.dsv3_fused_a_gemm_op(hidden_states, weight.t(), None, None) + equal = (out == shared).float().mean().item() + print(f" {equal:.4%} of the outputs bitwise equal to dsv3_fused_a_gemm_op") + assert equal >= 0.99 + assert_within_ulp(out, shared, element_ulp=4.0, rms_ulp=2.0) + + +def test_cuda_graph_replay() -> None: + """The decode step runs as one CUDA graph: the tensor maps and the list are fixed at capture, + the contents of the captured activations are not.""" + g = torch.Generator(device=DEV).manual_seed(3) + hidden_states = _tensor((8, HIDDEN_SIZE), BF16, g) + weight = _tensor((INDEXER_OUT_FEATURES, HIDDEN_SIZE), BF16, g) + l2_prefetch = [_tensor((4096, 2048), BF16, g), _tensor((4, 256, 512), BF16, g)] + op(hidden_states, weight, l2_prefetch) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(hidden_states, weight, l2_prefetch) + for _ in range(2): + hidden_states.copy_(_tensor((8, HIDDEN_SIZE), BF16, g)) + graph.replay() + torch.cuda.synchronize() + _compare(out, op.reference(hidden_states, weight, l2_prefetch)) + + +# ── the op's own rejections ─────────────────────────────────────────────── + + +def test_op_rejects_what_the_kernel_cannot_run() -> None: + """Loud on the op's own, with no guard interposed.""" + g = torch.Generator(device=DEV).manual_seed(5) + hidden_states = _tensor((8, HIDDEN_SIZE), BF16, g) + weight = _tensor((OUT_FEATURES, HIDDEN_SIZE), BF16, g) + hidden_7168, weight_7168 = _tensor((8, 7168), BF16, g), _tensor((2112, 7168), BF16, g) + strided_hidden_states = _tensor((HIDDEN_SIZE, 8), BF16, g).t() + strided_weight = _tensor((HIDDEN_SIZE, OUT_FEATURES), BF16, g).t() + flat = torch.zeros(8 * HIDDEN_SIZE + 8, dtype=BF16, device=DEV) + misaligned = flat[1 : 1 + 8 * HIDDEN_SIZE].view(8, HIDDEN_SIZE) + cases = [ + ("fp16 hidden_states", lambda: op(hidden_states.half(), weight, [])), + ("fp32 weight", lambda: op(hidden_states, weight.float(), [])), + ("3D hidden_states", lambda: op(hidden_states.unsqueeze(0), weight, [])), + ("hidden 7168", lambda: op(hidden_7168, weight_7168, [])), + ("2640 rows", lambda: op(hidden_states, _tensor((2640, HIDDEN_SIZE), BF16, g), [])), + ("9 tokens", lambda: op(_tensor((9, HIDDEN_SIZE), BF16, g), weight, [])), + ("0 tokens", lambda: op(hidden_states[:0], weight, [])), + ("strided hidden_states", lambda: op(strided_hidden_states, weight, [])), + ("strided weight", lambda: op(hidden_states, strided_weight, [])), + ("misaligned hidden_states", lambda: op(misaligned, weight, [])), + ("cpu weight", lambda: op(hidden_states, weight.cpu(), [])), + ("five list tensors", lambda: op(hidden_states, weight, [weight] * 5)), + ("cpu list tensor", lambda: op(hidden_states, weight, [weight.cpu()])), + ("strided list tensor", lambda: op(hidden_states, weight, [weight.t()])), + ("misaligned list tensor", lambda: op(hidden_states, weight, [misaligned])), + ] + for tag, call in cases: + with pytest.raises((RuntimeError, NotImplementedError)): + call() + pytest.fail(f"{tag} was accepted") + + +# ── the guards ──────────────────────────────────────────────────────────── + + +def _ok(out_features=INDEXER_OUT_FEATURES, tokens=8): + g = torch.Generator(device=DEV).manual_seed(out_features + tokens) + return _tensor((tokens, HIDDEN_SIZE), BF16, g), _tensor((out_features, HIDDEN_SIZE), BF16, g) + + +def test_guard_refuses_another_hidden_size() -> None: + g = torch.Generator(device=DEV).manual_seed(1) + hidden_states, weight = _tensor((8, 7168), BF16, g), _tensor((2112, 7168), BF16, g) + with pytest.raises(AssertionError, match=r"\[tokens, 6144\]"): + with validating(op): + op(hidden_states, weight, []) + + +def test_guard_refuses_a_weight_without_an_instance() -> None: + hidden_states, _ = _ok() + g = torch.Generator(device=DEV).manual_seed(2) + with pytest.raises(AssertionError, match="2624 or 2784 rows"): + with validating(op): + op(hidden_states, _tensor((2640, HIDDEN_SIZE), BF16, g), []) + + +@pytest.mark.parametrize("tokens", [0, 9, 16]) +def test_guard_refuses_a_token_count_outside_one_tile(tokens) -> None: + hidden_states, weight = _ok(tokens=max(tokens, 1)) + hidden_states = hidden_states[:tokens] + with pytest.raises(AssertionError, match="1 to 8 tokens"): + with validating(op): + op(hidden_states, weight, []) + + +@pytest.mark.parametrize("out_features", [OUT_FEATURES, INDEXER_OUT_FEATURES]) +def test_guard_refuses_fewer_sms_than_ctas(out_features, monkeypatch) -> None: + """The op runs anyway, in two waves; the caller takes the shared op there.""" + hidden_states, weight = _ok(out_features) + ctas = out_features // 16 + + class _Props: + multi_processor_count = ctas - 1 + + monkeypatch.setattr(torch.cuda, "get_device_properties", lambda device=None: _Props()) + with pytest.raises(AssertionError, match=f"{ctas} CTAs need as many SMs"): + with validating(op): + op(hidden_states, weight, []) From c78339ddd8eb8c93388d61fc73ed47871ab879b6 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 15:54:15 -0700 Subject: [PATCH 11/53] Add the GLM-5 router GEMV as the glm_router_gemm op kernels/glmMoeDsa/glmRouterGemm.cu computes the FP32 router logits of GLM-5's 256 experts at hidden 6144 for 1 to 16 tokens, one CTA per expert, with the loads and FMA order of the shared DeepSeek-V3 router GEMV, so the logits are bitwise those of dsv3_router_gemm_op. Before its grid-dependency wait each CTA prefetches its 12 KB weight row into L2, and it triggers its dependents right after the wait, before any logit is written: a kernel that follows it on the stream waits on its own grid dependency before it reads the logits. torch.ops.trtllm.glm_router_gemm(hidden_states, weight) takes the gate weight as stored and raises on shapes the kernel does not run. The modeling_v2 entry ops/gemm/router_gemm.py of the glm_5_2_nvfp4 sm_107 tp16 target states the shape gates in is_valid; its sm_107 test drives the cells against an FP64 reference and checks bitwise agreement with dsv3_router_gemm_op at every token count, CUDA graph replay and the op's rejections. Signed-off-by: Dongfeng Yu --- .../kernels/glmMoeDsa/glmRouterGemm.cu | 244 ++++++++++++++++++ .../kernels/glmMoeDsa/glmRouterGemm.h | 50 ++++ cpp/tensorrt_llm/thop/CMakeLists.txt | 2 +- cpp/tensorrt_llm/thop/glmRouterGemmOp.cpp | 94 +++++++ .../ops/gemm/router_gemm.py | 85 ++++++ .../gemm/test_modeling_v2_router_gemm.py | 173 +++++++++++++ 6 files changed, 647 insertions(+), 1 deletion(-) create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.cu create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.h create mode 100644 cpp/tensorrt_llm/thop/glmRouterGemmOp.cpp create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/router_gemm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_router_gemm.py diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.cu new file mode 100644 index 000000000000..c6196047a3b2 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.cu @@ -0,0 +1,244 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +namespace +{ + +// Convert 8 bfloat16 values from a uint4 to float array - optimized conversion +template +__device__ __forceinline__ void bf16_uint4_to_float8(uint4 const& vec, float* dst) +{ + __nv_bfloat16* bf16_ptr = reinterpret_cast<__nv_bfloat16*>(const_cast(&vec)); + +#pragma unroll + for (int i = 0; i < VPT; i++) + { + dst[i] = __bfloat162float(bf16_ptr[i]); + } +} + +// mat_a: the activations [kNumTokens, kHiddenDim]; mat_b: the weight [kNumExperts, kHiddenDim]; out: the logits +// [kNumTokens, kNumExperts]. Before the grid-dependency wait each CTA prefetches its weight row into L2 (no kernel of +// the step writes the weight), so the weight loads after the wait hit L2; the dependent grid may launch right after the +// wait instead of after the logits are written, and must wait on its own grid dependency before reading them. +template +__global__ __launch_bounds__(128, 1) void router_gemm_kernel(float* out, T const* mat_a, T const* mat_b) +{ + // Each block handles one expert column + int const n_idx = blockIdx.x; + int const tid = threadIdx.x; + constexpr int kWarpSize = 32; + constexpr int kNumWarps = kBlockSize / kWarpSize; + // Constants for this kernel + constexpr int k_elems_per_k_iteration = VPT * kBlockSize; + constexpr int k_iterations = kHiddenDim / k_elems_per_k_iteration; // Total K iterations + + // Initialize accumulators for all M rows + float acc[kNumTokens] = {}; + + // Shared memory for warp-level reduction + __shared__ float sm_reduction[kNumTokens][kNumWarps]; // kNumWarps + + // The weight row of expert n_idx (a column of B in column-major order) + T const* b_col = mat_b + n_idx * kHiddenDim; + + // Pre-compute k_base values for each iteration to help compiler optimize + int k_bases[k_iterations]; +#pragma unroll + for (int ki = 0; ki < k_iterations; ki++) + { + k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT; + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + { + // The primary grid does not write the weights, so this CTA's column can stream into L2 while the primary + // runs; the weight loads below then hit L2. What is loaded and the accumulation order are unchanged. + constexpr int kLineBytes = 128; + constexpr int kColumnBytes = kHiddenDim * static_cast(sizeof(T)); + static_assert(kColumnBytes % kLineBytes == 0, "the weight column must span whole 128-byte lines"); + char const* b_col_bytes = reinterpret_cast(b_col); + for (int offset = tid * kLineBytes; offset < kColumnBytes; offset += kBlockSize * kLineBytes) + { + asm volatile("prefetch.global.L2 [%0];" ::"l"(b_col_bytes + offset) : "memory"); + } + } + cudaGridDependencySynchronize(); + // Every dependent waits on its own grid dependency before it reads the logits, so it may launch now and overlap + // its prologue with this GEMV. + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + // Process the GEMM in chunks + for (int ki = 0; ki < k_iterations; ki++) + { + int const k_base = k_bases[ki]; + + // Load B matrix values using vector load (8 bf16 values) + uint4 b_vec = *reinterpret_cast(b_col + k_base); + + // Convert B values to float + float b_float[VPT]; + bf16_uint4_to_float8(b_vec, b_float); + +// Process each token +#pragma unroll + for (int m_idx = 0; m_idx < kNumTokens; m_idx++) + { + // Load both rows of A matrix using vector loads + uint4 a_vec = *reinterpret_cast(mat_a + (m_idx * kHiddenDim) + k_base); + + // Convert A values to float + float a_float[VPT]; + bf16_uint4_to_float8(a_vec, a_float); + +// Process elements in this chunk +#pragma unroll + for (int k = 0; k < VPT; k++) + { + float a = a_float[k]; + float b = b_float[k]; + acc[m_idx] += a * b; + } + } + } + + // Perform warp-level reduction + int const warpSize = 32; + int const warpId = tid / warpSize; + int const laneId = tid % warpSize; + + // Register for warp-level reduction results + float warp_result[kNumTokens]; + +#pragma unroll + for (int m_idx = 0; m_idx < kNumTokens; m_idx++) + { + warp_result[m_idx] = acc[m_idx]; + } + +// Perform warp-level reduction using optimized butterfly pattern +#pragma unroll + for (int m = 0; m < kNumTokens; m++) + { + float sum = warp_result[m]; + + // Butterfly reduction pattern + sum += __shfl_xor_sync(0xffffffff, sum, 16); + sum += __shfl_xor_sync(0xffffffff, sum, 8); + sum += __shfl_xor_sync(0xffffffff, sum, 4); + sum += __shfl_xor_sync(0xffffffff, sum, 2); + sum += __shfl_xor_sync(0xffffffff, sum, 1); + + // Only the first thread in each warp stores to shared memory + if (laneId == 0) + { + sm_reduction[m][warpId] = sum; + } + } + + __syncthreads(); + + // Final reduction across warps (only first thread) + if (tid == 0) + { +#pragma unroll + for (int m = 0; m < kNumTokens; m++) + { + float final_sum = 0.0f; + +// Sum across the kNumWarps +#pragma unroll + for (int w = 0; w < kNumWarps; w++) + { + final_sum += sm_reduction[m][w]; + } + + // Write final result + out[m * kNumExperts + n_idx] = final_sum; + } + } +} + +template +void launchRouterGemm( + float* logits, __nv_bfloat16 const* hiddenStates, __nv_bfloat16 const* weight, cudaStream_t stream) +{ + using T = __nv_bfloat16; + constexpr int VPT = 16 / sizeof(T); + constexpr int kBlockSize = 128; + cudaLaunchConfig_t config; + config.gridDim = kRouterNumExperts; + config.blockDim = kBlockSize; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = attrs; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, + router_gemm_kernel, logits, + hiddenStates, weight)); +} + +// Launches the instance of numTokens, one of kNumTokens..kRouterMaxTokens. +template +void dispatchRouterGemm( + float* logits, __nv_bfloat16 const* hiddenStates, __nv_bfloat16 const* weight, int numTokens, cudaStream_t stream) +{ + if constexpr (kNumTokens <= kRouterMaxTokens) + { + if (numTokens == kNumTokens) + { + launchRouterGemm(logits, hiddenStates, weight, stream); + } + else + { + dispatchRouterGemm(logits, hiddenStates, weight, numTokens, stream); + } + } + else + { + TLLM_THROW("The GLM router GEMV takes 1 to %d tokens, got %d", kRouterMaxTokens, numTokens); + } +} + +} // namespace + +void invokeRouterGemm( + float* logits, __nv_bfloat16 const* hiddenStates, __nv_bfloat16 const* weight, int numTokens, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO( + numTokens >= 1, "The GLM router GEMV takes 1 to %d tokens, got %d", kRouterMaxTokens, numTokens); + dispatchRouterGemm<1>(logits, hiddenStates, weight, numTokens, stream); +} + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.h new file mode 100644 index 000000000000..c103c7eeb7e7 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.h @@ -0,0 +1,50 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +//! The router GEMV's K: the GLM-5 hidden size. +constexpr int kRouterHiddenSize = 6144; +//! Routed experts: one CTA computes one expert's logits. +constexpr int kRouterNumExperts = 256; +//! Tokens per call at most: one kernel instance per token count. +constexpr int kRouterMaxTokens = 16; + +//! logits[t, e] = sum_k hiddenStates[t, k] * weight[e, k]: BF16 operands, FP32 accumulation and FP32 logits, all +//! dense row-major. hiddenStates is [numTokens, 6144] with numTokens in [1, 16], weight [256, 6144] and logits +//! [numTokens, 256]. +//! +//! Launched with the programmatic-dependent-launch attribute when TRTLLM_ENABLE_PDL is set. Before its grid-dependency +//! wait each CTA prefetches its weight row into L2; right after the wait, before any logit is written, it triggers its +//! dependents. Every kernel that follows on the stream must therefore wait on its own grid dependency before it reads +//! the logits or writes hiddenStates. +void invokeRouterGemm( + float* logits, __nv_bfloat16 const* hiddenStates, __nv_bfloat16 const* weight, int numTokens, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index faaa106a4995..4ebb89865614 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -120,7 +120,7 @@ add_th_common_family( # because moe/utils/moeUtilOp.cpp depends on OSS-specific MoE utility types and # function signatures. add_th_common_family(th_common_moe fusedTopkSoftmax.cpp llama4MinLatency.cpp - noAuxTcOp.cpp dsv3RouterGemmOp.cpp) + noAuxTcOp.cpp dsv3RouterGemmOp.cpp glmRouterGemmOp.cpp) # MoE torch ops live in moe//; that directory adds them to # th_common_moe with target_sources() so the list stays next to the files. diff --git a/cpp/tensorrt_llm/thop/glmRouterGemmOp.cpp b/cpp/tensorrt_llm/thop/glmRouterGemmOp.cpp new file mode 100644 index 000000000000..cbaa789904ce --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmRouterGemmOp.cpp @@ -0,0 +1,94 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.h" + +#include +#include +#include +#include + +namespace tk = tensorrt_llm::kernels; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +// The kernel reads both operands in 16-byte vectors. +constexpr int64_t kVectorBytes = 16; + +bool isAligned(torch::Tensor const& tensor) +{ + return reinterpret_cast(tensor.data_ptr()) % kVectorBytes == 0; +} + +} // namespace + +//! The GLM-5 MoE router logits for up to 16 tokens. +//! +//! @param hiddenStates BF16 CUDA tensor [num_tokens, 6144], contiguous, num_tokens in [1, 16]. +//! @param weight BF16 CUDA tensor [256, 6144], contiguous: the router (gate) weight as stored. +//! @return FP32 tensor [num_tokens, 256] = hiddenStates @ weight^T, FP32 accumulation. +torch::Tensor glmRouterGemm(torch::Tensor const& hiddenStates, torch::Tensor const& weight) +{ + namespace glm = tk::glm_moe_dsa; + TORCH_CHECK(hiddenStates.is_cuda() && weight.is_cuda(), "glm_router_gemm expects CUDA tensors"); + TORCH_CHECK(hiddenStates.device() == weight.device(), "hidden_states and weight must be on the same device"); + TORCH_CHECK(hiddenStates.scalar_type() == torch::kBFloat16 && weight.scalar_type() == torch::kBFloat16, + "glm_router_gemm expects BF16 hidden_states and weight, got ", hiddenStates.scalar_type(), " and ", + weight.scalar_type()); + TORCH_CHECK(hiddenStates.dim() == 2 && weight.dim() == 2, + "glm_router_gemm expects hidden_states [num_tokens, 6144] and weight [256, 6144]"); + TORCH_CHECK(hiddenStates.size(1) == glm::kRouterHiddenSize && weight.size(1) == glm::kRouterHiddenSize, + "glm_router_gemm serves hidden size ", glm::kRouterHiddenSize, ", got hidden_states ", hiddenStates.sizes(), + " and weight ", weight.sizes()); + TORCH_CHECK(weight.size(0) == glm::kRouterNumExperts, "glm_router_gemm serves ", glm::kRouterNumExperts, + " experts, got ", weight.size(0)); + TORCH_CHECK(hiddenStates.size(0) >= 1 && hiddenStates.size(0) <= glm::kRouterMaxTokens, + "glm_router_gemm takes 1 to ", glm::kRouterMaxTokens, " tokens, got ", hiddenStates.size(0)); + TORCH_CHECK(hiddenStates.is_contiguous() && weight.is_contiguous(), + "glm_router_gemm expects contiguous row-major hidden_states and weight"); + TORCH_CHECK(isAligned(hiddenStates) && isAligned(weight), + "glm_router_gemm expects 16-byte aligned hidden_states and weight"); + + c10::cuda::CUDAGuard const deviceGuard{hiddenStates.device()}; + int const numTokens = static_cast(hiddenStates.size(0)); + torch::Tensor logits + = torch::empty({numTokens, glm::kRouterNumExperts}, hiddenStates.options().dtype(torch::kFloat32)); + auto const stream = at::cuda::getCurrentCUDAStream(hiddenStates.get_device()); + glm::invokeRouterGemm(logits.data_ptr(), reinterpret_cast<__nv_bfloat16 const*>(hiddenStates.data_ptr()), + reinterpret_cast<__nv_bfloat16 const*>(weight.data_ptr()), numTokens, stream); + return logits; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("glm_router_gemm(Tensor hidden_states, Tensor weight) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_router_gemm", &tensorrt_llm::torch_ext::glmRouterGemm); +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/router_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/router_gemm.py new file mode 100644 index 000000000000..990fe197a7ea --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/router_gemm.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GLM-5 MoE router logits (256 experts at hidden 6144, FP32 from BF16) for up to 16 tokens, via +the trtllm glm_router_gemm op.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +HIDDEN_SIZE = 6144 +NUM_EXPERTS = 256 +MAX_TOKENS = 16 + + +class GlmRouterGemm(OpWrapper): + """`hidden_states @ weight.T` in FP32 from BF16 operands: the router logits `[tokens, 256]`. + + `weight` is the gate weight as stored, `[256, 6144]` BF16 row-major. One CTA per expert. Before + its grid-dependency wait each CTA prefetches its weight row into L2; right after the wait, + before any logit is written, it lets its dependents launch. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The c1 decode step routes the MTP layer's BF16 MoE here once per draft step (7 calls); the + # 75 MoE layers of the target compute their router logits inside the routing kernel. + CELLS: tuple[Cell, ...] = ( + Cell( + why="MTP draft steps: one row (draft step 0 after its row select, steps 1-6)", + spec=dict(tokens=1), + ), + Cell( + why="8 rows: a verify-width step routed without the row select", + spec=dict(tokens=8), + ), + Cell(why="16 tokens: the op's ceiling, its largest instance", spec=dict(tokens=16)), + ) + + note = """ + PDL (`TRTLLM_ENABLE_PDL`): the kernel reads only `weight` before its grid-dependency wait (an + L2 prefetch), so no kernel launched earlier in the step may write it. Its dependents launch + as soon as it passes the wait, before the logits exist: each must wait on its own grid + dependency before it reads the logits or writes `hidden_states`. At c1 the dependent is the + low-M BF16 MoE's FC1 kernel, which waits before it routes; its FC2 kernel reads the logits + before its own wait, which holds because FC1 releases it only after FC1's wait. + + Preconditions, enforced loudly by the op: CUDA BF16 2-D `hidden_states` and `weight` on one + device, contiguous and 16-byte aligned; hidden 6144; 256 experts; 1 to 16 tokens. + + Numerics: per thread an FP32 FMA chain over its 48 products, then a warp butterfly and the four + warp sums in order; deterministic, and bitwise the logits of the shared + `dsv3_router_gemm_op`, whose kernel sums in the same order. + """ + + def __call__(self, hidden_states: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + return torch.ops.trtllm.glm_router_gemm(hidden_states, weight) + + def reference(self, hidden_states: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """`hidden_states @ weight.T` accumulated in FP64 and rounded to FP32 once. + + The op's output is FP32, so an FP32-accumulated reference would carry rounding of the + same size as the kernel's; FP64 keeps the reference the more accurate side. + """ + return (hidden_states.double() @ weight.double().t()).float() + + def is_valid(self, hidden_states: torch.Tensor, weight: torch.Tensor) -> None: + # The shape gates of the kernel, which a target checks to route anything else to the + # shared dsv3_router_gemm_op. The op raises on them as well. + assert hidden_states.dim() == 2 and hidden_states.shape[1] == HIDDEN_SIZE, ( + f"hidden_states must be [tokens, {HIDDEN_SIZE}]" + ) + # One CTA per expert row of the [256, 6144] gate weight. + assert tuple(weight.shape) == (NUM_EXPERTS, HIDDEN_SIZE), ( + f"weight must be [{NUM_EXPERTS}, {HIDDEN_SIZE}], got {tuple(weight.shape)}" + ) + # One kernel instance per token count. + assert 1 <= hidden_states.shape[0] <= MAX_TOKENS, ( + f"the kernel takes 1 to {MAX_TOKENS} tokens, got {hidden_states.shape[0]}" + ) + + +router_gemm = GlmRouterGemm() diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_router_gemm.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_router_gemm.py new file mode 100644 index 000000000000..29a981ee05a2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_router_gemm.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM router_gemm entry. + +The entry owns the contract. This file turns a cell into tensors and holds the claims no cell +carries: bitwise agreement with the shared router GEMV, determinism, every token count up to the +ceiling, CUDA graph replay, the op's own rejections, and the guards. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.router_gemm import ( # noqa: E501 + HIDDEN_SIZE, + MAX_TOKENS, + NUM_EXPERTS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.router_gemm import ( # noqa: E501 + router_gemm as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402,F401 — needs the path declared above; fixture by name + ULP, + assert_within_ulp, + certifying, + validating, +) + +assert torch.cuda.is_available(), "router_gemm requires a CUDA device" + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM router GEMV is certified on sm_107 only" +) + +DEV = "cuda" +BF16 = torch.bfloat16 + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """FP32 accumulation against an FP64 reference, scored in FP32 ulps. + + Both sides sum the same products (a BF16 x BF16 product is exact in FP32); the reference in + FP64, the kernel in FP32: per thread a chain of 48 FMAs, then a 5-level warp butterfly and the + four warp sums. With randn operands the logits sit near sqrt(6144) ~ 80 (row maxima ~230), + and the FP32 roundings of those additions add up to about 2.5e-5 per logit: about 3 ulp + (2^-23) of relative RMS and about 3 ulp of the row maximum at worst. + + So: 64 ulp element-wise against the row scale and 32 ulp of relative RMS, tenfold over the + estimate. A BF16 rounding of the logits alone would be 2^-9, four orders above, and a missing + k chunk, expert or token misses by more. + """ + assert_within_ulp(out, ref, element_ulp=64.0, rms_ulp=32.0, ulp=ULP[torch.float32]) + + +def _build(tokens, seed): + """randn activations and gate weight, both BF16.""" + g = torch.Generator(device=DEV).manual_seed(seed) + hidden_states = torch.randn(tokens, HIDDEN_SIZE, generator=g, device=DEV).to(BF16) + weight = torch.randn(NUM_EXPERTS, HIDDEN_SIZE, generator=g, device=DEV).to(BF16) + return hidden_states, weight + + +# ── the certified cells ─────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture + certifying(op) + hidden_states, weight = _build(cell.spec["tokens"], zlib.crc32(cell.why.encode()) % 2**31) + logits = op(hidden_states, weight) + assert logits.shape == (hidden_states.shape[0], NUM_EXPERTS) + assert logits.dtype == torch.float32 and logits.is_contiguous() + _compare(logits, op.reference(hidden_states, weight)) + + +# ── the claims no cell carries ──────────────────────────────────────────── + + +@pytest.mark.parametrize("tokens", range(1, MAX_TOKENS + 1)) +def test_bitwise_the_shared_router_gemv(tokens) -> None: + """dsv3_router_gemm_op runs the same loads and FMA chain; only the trigger and the L2 + prefetch differ, and neither touches the arithmetic. Every token count is an instance.""" + hidden_states, weight = _build(tokens, seed=tokens) + logits = op(hidden_states, weight) + shared = torch.ops.trtllm.dsv3_router_gemm_op(hidden_states, weight.t(), None, torch.float32) + assert torch.equal(logits, shared) + _compare(logits, op.reference(hidden_states, weight)) + + +def test_deterministic() -> None: + hidden_states, weight = _build(8, seed=21) + first = op(hidden_states, weight) + second = op(hidden_states, weight) + assert torch.equal(first, second) and first.data_ptr() != second.data_ptr() + + +def test_cuda_graph_replay() -> None: + """The decode step runs as one CUDA graph: the captured activations change between replays.""" + hidden_states, weight = _build(1, seed=4) + op(hidden_states, weight) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + logits = op(hidden_states, weight) + for seed in (5, 6): + hidden_states.copy_(_build(1, seed)[0]) + graph.replay() + torch.cuda.synchronize() + _compare(logits, op.reference(hidden_states, weight)) + + +# ── the op's own rejections ─────────────────────────────────────────────── + + +def test_op_rejects_what_the_kernel_cannot_run() -> None: + """Loud on the op's own, with no guard interposed.""" + hidden_states, weight = _build(8, seed=9) + g = torch.Generator(device=DEV).manual_seed(10) + hidden_7168 = torch.randn(8, 7168, generator=g, device=DEV).to(BF16) + weight_7168 = torch.randn(NUM_EXPERTS, 7168, generator=g, device=DEV).to(BF16) + weight_128 = weight[:128] + hidden_17 = torch.randn(17, HIDDEN_SIZE, generator=g, device=DEV).to(BF16) + strided_hidden_states = torch.randn(HIDDEN_SIZE, 8, generator=g, device=DEV).to(BF16).t() + flat = torch.zeros(8 * HIDDEN_SIZE + 8, dtype=BF16, device=DEV) + misaligned = flat[1 : 1 + 8 * HIDDEN_SIZE].view(8, HIDDEN_SIZE) + cases = [ + ("fp16 hidden_states", lambda: op(hidden_states.half(), weight)), + ("fp32 weight", lambda: op(hidden_states, weight.float())), + ("3D hidden_states", lambda: op(hidden_states.unsqueeze(0), weight)), + ("hidden 7168", lambda: op(hidden_7168, weight_7168)), + ("128 experts", lambda: op(hidden_states, weight_128)), + ("17 tokens", lambda: op(hidden_17, weight)), + ("0 tokens", lambda: op(hidden_states[:0], weight)), + ("strided hidden_states", lambda: op(strided_hidden_states, weight)), + ("weight.t()", lambda: op(hidden_states, weight.t())), + ("misaligned hidden_states", lambda: op(misaligned, weight)), + ("cpu weight", lambda: op(hidden_states, weight.cpu())), + ] + for tag, call in cases: + with pytest.raises((RuntimeError, NotImplementedError)): + call() + pytest.fail(f"{tag} was accepted") + + +# ── the guards ──────────────────────────────────────────────────────────── + + +def test_guard_refuses_another_hidden_size() -> None: + g = torch.Generator(device=DEV).manual_seed(1) + hidden_states = torch.randn(1, 7168, generator=g, device=DEV).to(BF16) + weight = torch.randn(NUM_EXPERTS, 7168, generator=g, device=DEV).to(BF16) + with pytest.raises(AssertionError, match=r"\[tokens, 6144\]"): + with validating(op): + op(hidden_states, weight) + + +def test_guard_refuses_another_expert_count() -> None: + hidden_states, weight = _build(1, seed=2) + with pytest.raises(AssertionError, match=r"weight must be \[256, 6144\]"): + with validating(op): + op(hidden_states, weight[:128]) + + +@pytest.mark.parametrize("tokens", [0, 17, 32]) +def test_guard_refuses_a_token_count_without_an_instance(tokens) -> None: + hidden_states, weight = _build(max(tokens, 1), seed=3) + with pytest.raises(AssertionError, match="1 to 16 tokens"): + with validating(op): + op(hidden_states[:tokens], weight) From f5f79e6e74713bc30ed6be4b8e7e2df0233ef937 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:02:06 -0700 Subject: [PATCH 12/53] Run the GLM-5.2 decode step's attention as target code with the fused-A GEMM The decode step runs MLA with the DSA indexer itself: q_a, the compressed KV and its rope part come from the GLM fused-A GEMM over the stored kv_a_proj_with_mqa weight for up to 8 rows, and the built-in norms, q_b, indexer projection, K-cache update, sparse generation attention and o_proj compute the rest in the order the built-in DSA path runs them. Signed-off-by: Dongfeng Yu --- .../glm_5_2_nvfp4__sm_107__tp16/modeling.py | 103 +++++++++++++++--- 1 file changed, 88 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py index 4b356ab23023..b926bb4f2939 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py @@ -20,16 +20,23 @@ from transformers import PretrainedConfig from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata +from tensorrt_llm._torch.attention.backends.sparse.dsa.module import ( + forward_generation_sparse_attn, +) from tensorrt_llm._torch.distributed import AllReduceFusionOp, AllReduceParams from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3ForCausalLM, DeepseekV3Model from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel from tensorrt_llm._torch.speculative import SpecMetadata from . import weights as _weights from .ops.embedding.fused_embedding import fused_embedding from .ops.gemm.dense_mlp_small_m import INTER_PER_RANK, dense_mlp_small_m +from .ops.gemm.fused_a_gemm import MAX_TOKENS as FUSED_A_MAX_TOKENS +from .ops.gemm.fused_a_gemm import OUT_FEATURES as FUSED_A_OUT_FEATURES +from .ops.gemm.fused_a_gemm import fused_a_gemm # The deployment this target is written for; route() selects it on the same # facts, and the constructor states them again for a direct construction. @@ -79,6 +86,11 @@ class DecodeTarget: following norm; the experts finalize before theirs (16 ranks span nodes). - The last all-reduce applies the model's final norm (``next_layer_layernorm``). + Attention is MLA with the DSA indexer on the generation rows: q_a, the + compressed KV and its rope part come from one GEMM over the stored + kv_a_proj_with_mqa weight (up to 8 rows), and the built-in modules compute + the rest. + The embedding table is whole on every rank, so the lookup is one local kernel. The dense MLP runs as two weight-streaming kernels on this rank's 768 intermediate columns for up to 8 rows (every c1 step: one request and @@ -88,7 +100,15 @@ class DecodeTarget: def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: self.core = core for layer in core.layers[: core.num_hidden_layers]: + mla = layer.self_attn assert layer.allreduce is not None and layer.next_layer_layernorm is not None + assert not mla.apply_rotary_emb and not mla.mapping.has_cp_helix() + assert tuple(mla.kv_a_proj_with_mqa.weight.shape) == ( + FUSED_A_OUT_FEATURES, + core.config.hidden_size, + ) and mla.kv_a_proj_with_mqa.weight.dtype == torch.bfloat16, ( + f"layer {layer.layer_idx}: kv_a_proj_with_mqa is not the BF16 [q_a; kv_a] weight" + ) assert layer.input_layernorm.nvfp4_scale is None, "kv_a_proj is not NVFP4" if isinstance(layer.mlp, GatedMLP): assert ( @@ -120,27 +140,78 @@ def forward( inputs_embeds = fused_embedding(input_ids, core.embed_tokens.weight) residual = inputs_embeds hidden_states = core.layers[0].input_layernorm(inputs_embeds) + rows = hidden_states.shape[0] small_m = torch.ops.trtllm.glm_dense_mlp_small_m_supported( - hidden_states.shape[0], core.config.hidden_size, INTER_PER_RANK + rows, core.config.hidden_size, INTER_PER_RANK ) + fused_a = rows <= FUSED_A_MAX_TOKENS for layer in core.layers[: core.num_hidden_layers]: if isinstance(layer.mlp, GatedMLP): hidden_states, residual = self._dense_layer( - layer, position_ids, hidden_states, residual, attn_metadata, small_m + layer, position_ids, hidden_states, residual, attn_metadata, fused_a, small_m ) else: hidden_states, residual = self._moe_layer( - layer, position_ids, hidden_states, residual, attn_metadata + layer, position_ids, hidden_states, residual, attn_metadata, fused_a ) return hidden_states @staticmethod - def _dense_layer(layer, position_ids, hidden_states, residual, attn_metadata, small_m): - hidden_states = layer.self_attn( - position_ids=position_ids, - hidden_states=hidden_states, - attn_metadata=attn_metadata, - all_reduce_params=AllReduceParams(enable_allreduce=True), + def _attention(layer, position_ids, hidden_states, attn_metadata, all_reduce_params, fused_a): + """MLA with the DSA indexer over generation rows, through o_proj.""" + mla = layer.self_attn + num_tokens = attn_metadata.num_tokens + if fused_a: + kv_a = fused_a_gemm(hidden_states, mla.kv_a_proj_with_mqa.weight, []) + else: + kv_a = mla.kv_a_proj_with_mqa(hidden_states) + q, compressed_kv, k_pe = kv_a.split( + [mla.q_lora_rank, mla.kv_lora_rank, mla.qk_rope_head_dim], -1 + ) + (q, qr), compressed_kv = maybe_execute_in_parallel( + lambda: mla._q_a_layernorm_maybe_fused(q, return_norm_out=True), + lambda: mla.kv_a_layernorm(compressed_kv), + mla.ln_events[0], + mla.ln_events[1], + mla.aux_stream, + ) + latent_cache = torch.concat([compressed_kv, k_pe], dim=-1) + q = mla.q_b_proj(q) + indexer = mla.mqa.indexer + indexer_intermediates = None + if indexer is not None: + indexer_intermediates = [ + t[:num_tokens] + for t in indexer.pre_indexer_proj(qr, hidden_states, position_ids) + ] + indexer._update_k_cache(indexer_intermediates[1], indexer_intermediates[2], attn_metadata) + output = hidden_states.new_empty([hidden_states.shape[0], mla.attention_output_hidden_size]) + forward_generation_sparse_attn( + mla, + q[:num_tokens], + compressed_kv[:num_tokens], + k_pe[:num_tokens], + attn_metadata, + output[:num_tokens], + position_ids=position_ids[..., :num_tokens], + latent_cache=latent_cache[:num_tokens], + indexer_intermediates=indexer_intermediates, + ) + if indexer is not None: + indexer.maybe_join_prev_topk_copy() + return mla.o_proj(output, all_reduce_params=all_reduce_params, layer_idx=mla.layer_idx) + + @staticmethod + def _dense_layer( + layer, position_ids, hidden_states, residual, attn_metadata, fused_a, small_m + ): + hidden_states = DecodeTarget._attention( + layer, + position_ids, + hidden_states, + attn_metadata, + AllReduceParams(enable_allreduce=True), + fused_a, ) hidden_states, residual = layer.post_attention_layernorm(hidden_states, residual) if small_m: @@ -163,12 +234,14 @@ def _dense_layer(layer, position_ids, hidden_states, residual, attn_metadata, sm ) @staticmethod - def _moe_layer(layer, position_ids, hidden_states, residual, attn_metadata): - hidden_states = layer.self_attn( - position_ids=position_ids, - hidden_states=hidden_states, - attn_metadata=attn_metadata, - all_reduce_params=AllReduceParams(enable_allreduce=False), + def _moe_layer(layer, position_ids, hidden_states, residual, attn_metadata, fused_a): + hidden_states = DecodeTarget._attention( + layer, + position_ids, + hidden_states, + attn_metadata, + AllReduceParams(enable_allreduce=False), + fused_a, ) hidden_states, residual = layer.allreduce( hidden_states, From 4da71b5125a53d49bdefe567b3045bf1d23aa511 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 15:59:32 -0700 Subject: [PATCH 13/53] Add the glm_moe_dsa tp16 target's low-M split-K GEMM ops The target carries its own copy of the low-M BF16 split-K CuTe DSL GEMM (ops/gemm/_low_m_splitk.py, adapted from FlashInfer's dense_bf16_gemm_sm100_splitk) with the occupancy policy that refines the source default tactic to a cluster of eight K slices, the Lamport publish epilogue and the in-kernel RMSNorm B producer with an FP32 output, plus the fused-epilogue kernels (ops/gemm/_low_m_splitk_fused.py): the row-RMS and FP8 RoPE q_b epilogue and the DSA indexer Q tail. Entries, one per call the decode step makes, all certified on sm_107: - low_m_gemm: x @ weight.T for the shared expert's gate_up and down, the LM head shard and MTP draft step 0's eh_proj. - q_b_fp8_rope: q_b with k_b absorbed and the q_a RMSNorm folded into the weight, writing the RoPE'd E4M3 attention Q; regular, beside-top-k and early-gather launch plans. - indexer_q_tail: the indexer wq_b from the raw q_a with the prologue's Q tail (RoPE, UE8M0 scales, E4M3 Q, scaled head weights) in the epilogue. - o_proj_publish: o_proj storing its BF16 output into this rank's slot of the post-attention MNNVL one-shot all-reduce's current Lamport buffer. - draft_hidden_half: this rank's FP32 share of the next draft step's eh_proj hidden half, hnorm computed in the kernel for one row. Each entry has a GPU test over its cells with a derived band, a CUDA graph replay check where the target captures it, and guard tests for the inputs is_valid refuses. Signed-off-by: Dongfeng Yu --- .../ops/gemm/_low_m_splitk.py | 1888 +++++++++++++++ .../ops/gemm/_low_m_splitk_fused.py | 2033 +++++++++++++++++ .../ops/gemm/draft_hidden_half.py | 136 ++ .../ops/gemm/indexer_q_tail.py | 253 ++ .../ops/gemm/low_m_gemm.py | 109 + .../ops/gemm/o_proj_publish.py | 174 ++ .../ops/gemm/q_b_fp8_rope.py | 213 ++ .../test_modeling_v2_draft_hidden_half.py | 148 ++ .../gemm/test_modeling_v2_indexer_q_tail.py | 147 ++ .../gemm/test_modeling_v2_low_m_gemm.py | 104 + .../gemm/test_modeling_v2_o_proj_publish.py | 158 ++ .../gemm/test_modeling_v2_q_b_fp8_rope.py | 192 ++ 12 files changed, 5555 insertions(+) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk_fused.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/draft_hidden_half.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/indexer_q_tail.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/low_m_gemm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/o_proj_publish.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/q_b_fp8_rope.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_draft_hidden_half.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_indexer_q_tail.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_low_m_gemm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_o_proj_publish.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_q_b_fp8_rope.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk.py new file mode 100644 index 000000000000..5613cb522764 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk.py @@ -0,0 +1,1888 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Adapted from FlashInfer (https://github.com/flashinfer-ai/flashinfer) under the Apache-2.0 +# license. Original file: flashinfer/gemm/kernels/dense_bf16_gemm_sm100_splitk.py +# Original copyright: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +"""SM10x low-M BF16/FP16 GEMM with an in-kernel cluster split-K reduction (CuTe DSL). + +Each cluster rank accumulates an exact K slice in FP32. Peers publish partials +to rank 0 through DSMEM; rank 0 reduces, casts, and stores once. The public +``A[M, K] @ B[K, N]`` problem is swapped internally, so tile dimensions below +use kernel coordinates: kernel-M carries public N and kernel-N carries public M. + +This target's own copy of the kernel. Beside the plain GEMM it carries what the +target's decode step needs and FlashInfer's kernel does not have: + +- ``splitk_default_tactic``: the source default tactic, refined to a cluster of + eight K slices where that fills otherwise idle SMs; +- ``LamportPublish`` (``run_splitk_publish``): the output stored straight into + this rank's slot of an MNNVL one-shot all-reduce workspace; +- ``HnormB`` and an FP32 output (``run_splitk_hidden_half``): B produced in the + kernel as FlashInfer's RMSNorm of one activation row, the result kept in FP32. + +The launchers validate their operands in Python (``ValueError``) and compile one +kernel per specialization on first use. +""" + +from __future__ import annotations + +import dataclasses +import functools +import struct + +import cuda.bindings.driver as _cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import torch as _torch +from cutlass import Float32, Int32, Int64 +from cutlass._mlir.dialects import llvm +from cutlass.cute import experimental as cute_ext +from cutlass.cute.nvgpu import tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +#: Per-CTA SMEM capacity reported by CuTeDSL on SM100/SM103. +_SMEM_CAPACITY_BYTES = utils.get_smem_capacity_in_bytes("sm_100") + +#: K extent of one CTA tile. +_CTA_K = 128 + +#: Kernel-M tiles; 64 increases CTA count for low-M decode shapes. +_SUPPORTED_MMA_M = (64, 128) + +#: Kernel-N carries public M, which is limited to 32. +_SUPPORTED_MMA_N = (8, 16, 32) + +#: Physical cluster-K sizes; split 1 compiles out the DSMEM path. +_SUPPORTED_SPLIT_K = (1, 2, 3, 4) + +#: Largest public M this low-M policy serves. +MAX_M = 32 + +#: Bytes per FP32 partial exchanged through DSMEM. +_FP32_BYTES = 4 + +#: DSMEM mailbox base alignment, in bytes. +_MAILBOX_ALIGN_BYTES = 128 + +#: Size and alignment of one mbarrier, in bytes. +_MBARRIER_BYTES = 8 + +#: Size and alignment of the TMEM base pointer slot. +_TMEM_POINTER_BYTES = 4 + +#: Bytes per BF16/FP16 element. +_AB_ELEMENT_BYTES = 2 + +#: Alignment of the A/B shared-memory buffers. +_AB_BUFFER_ALIGN_BYTES = 1024 + +#: A/B pipeline stage bounds. One stage is useful only for a single K tile. +_MIN_AB_STAGES = 1 +_DEFAULT_MIN_AB_STAGES = 2 +_MAX_AB_STAGES = 12 + + +@dataclasses.dataclass(frozen=True) +class SplitKTactic: + """One specialization; mma_m carries public N and mma_n carries public M.""" + + mma_m: int + mma_n: int + split_k: int + ab_stages: int + + +def _align_up(value: int, alignment: int) -> int: + return ((value + alignment - 1) // alignment) * alignment + + +def _smem_bytes( + tactic: SplitKTactic, + ab_stages: int, +) -> int: + """Mirror the device allocator's shared-memory layout.""" + cursor = ( + _align_up( + tactic.mma_m * _CTA_K * _AB_ELEMENT_BYTES * ab_stages, + _AB_BUFFER_ALIGN_BYTES, + ) + + tactic.mma_n * _CTA_K * _AB_ELEMENT_BYTES * ab_stages + ) + + cursor = _align_up(cursor, _MBARRIER_BYTES) + cursor += 2 * ab_stages * _MBARRIER_BYTES + cursor += 3 * _MBARRIER_BYTES + cursor = _align_up(cursor, _TMEM_POINTER_BYTES) + cursor += _TMEM_POINTER_BYTES + + if tactic.split_k == 1: + return cursor + + return ( + _align_up( + _align_up(cursor, _MAILBOX_ALIGN_BYTES) + + (tactic.split_k - 1) * tactic.mma_m * tactic.mma_n * _FP32_BYTES, + _MBARRIER_BYTES, + ) + + _MBARRIER_BYTES + ) + + +#: Headroom over the mirrored layout when a launch reserves only the tactic's footprint. +_LAUNCH_SMEM_HEADROOM_BYTES = 2048 + +#: Alignment of the Lamport-publish staging tile. +_PUBLISH_ALIGN_BYTES = 16 + +#: Named barrier of the four epilogue warps in the Lamport-publish epilogue (0 is the CTA-wide one). +_PUBLISH_BARRIER_ID = 1 + +#: BF16 elements per 16-byte Lamport-publish store. +_PUBLISH_VEC_ELEMENTS = 8 + + +def _publish_smem_bytes(tactic: SplitKTactic, ab_stages: int) -> int: + """Shared memory of the Lamport-publish kernel: the base layout, then the FP32 output tile.""" + return ( + _align_up(_smem_bytes(tactic, ab_stages), _PUBLISH_ALIGN_BYTES) + + tactic.mma_m * tactic.mma_n * _FP32_BYTES + ) + + +@dataclasses.dataclass(frozen=True) +class LamportPublish: + """Publish the GEMM output into the current buffer of an MNNVL one-shot all-reduce workspace + instead of writing C. + + The owner CTA of each output tile stores its BF16 values, -0.0 written as +0.0, through the + workspace's multicast address at [token][rank][column] of the buffer the workspace flags mark + current: the shard store of the one-shot kernel, which then only polls. Every field is a + compile-time constant of the kernel. + + ``trigger_after_wait``: the PDL dependents (the all-reduce) are released by the epilogue warps + right after the grid dependency wait instead of by the weight-load warp once its loads are + issued, so that the all-reduce's reads before its own wait (the flags, the residual) follow + the completion of this GEMM's primary. + """ + + world_size: int + rank: int + trigger_after_wait: bool = True + + +#: Row length whose FlashInfer RMSNorm reduction ``HnormB`` reproduces. +HNORM_B_HIDDEN = 6144 + +#: BF16 elements per 16-byte vector of the HnormB producer. +_HNORM_B_VEC_ELEMENTS = 8 + +#: Lanes of a warp. +_WARP_LANES = 32 + + +@dataclasses.dataclass(frozen=True) +class HnormB: + """Produce B in the kernel instead of loading it: B is FlashInfer's RMSNorm + (``flashinfer.norm.rmsnorm``, its CuTe DSL kernel at H = 6144) of the one activation row + passed as B, with the weight passed as the bias operand, rounded to BF16. Warp 1 computes the + norm with that kernel's arithmetic and summation order and writes this CTA's K slice into the + B stages, so the MMA reads what the TMA would have loaded from the norm's output and the GEMM's + result is bit-identical to the GEMM of that output. Every field is a compile-time constant. + """ + + hidden: int + eps: float + + +def hnorm_b_supported(tactic: SplitKTactic, hidden: int) -> bool: + """Whether ``HnormB`` serves ``tactic`` for a row of ``hidden``: the reproduced width, 8-row B + tiles (one K_SW128 swizzle atom row group per 64 columns), every K tile of a CTA in its own + stage, and a CTA slice of whole 32-vector groups (one warp, one vector per lane per group).""" + if hidden != HNORM_B_HIDDEN or tactic.mma_n != 8: + return False + if hidden % (_CTA_K * tactic.split_k): + return False + k_tiles = hidden // _CTA_K // tactic.split_k + return ( + k_tiles <= tactic.ab_stages + and (k_tiles * _CTA_K // _HNORM_B_VEC_ELEMENTS) % _WARP_LANES == 0 + ) + + +def tactic_launch_smem_bytes(tactic: SplitKTactic) -> int: + """Dynamic shared memory for a launch that reserves only what ``tactic`` uses. + + The default launch reserves the full capacity, which keeps one CTA per SM. This size + lets other CTAs share the SM: two CTAs of the c1 q_b tactic (128, 8, 8, 2) fit on one. + """ + return min( + _align_up(_smem_bytes(tactic, tactic.ab_stages), 1024) + _LAUNCH_SMEM_HEADROOM_BYTES, + _SMEM_CAPACITY_BYTES, + ) + + +def _max_ab_stages_for( + tactic: SplitKTactic, + smem_capacity: int, +) -> int: + return next( + ( + stages + for stages in range(_MAX_AB_STAGES, -1, -1) + if _smem_bytes(tactic, stages) <= smem_capacity + ), + 0, + ) + + +def validate_tactic( + tactic: SplitKTactic, + m: int, + n: int, + k: int, + *, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> None: + """Reject a tactic that cannot serve ``(m, n, k)``.""" + if tactic.mma_m not in _SUPPORTED_MMA_M: + raise ValueError(f"unsupported mma_m={tactic.mma_m}") + if tactic.mma_n not in _SUPPORTED_MMA_N: + raise ValueError(f"unsupported mma_n={tactic.mma_n}") + # Split 8 (splitk_default_tactic) is served beside the cataloged cluster sizes. + if tactic.split_k not in _SUPPORTED_SPLIT_K and tactic.split_k != 8: + raise ValueError(f"unsupported split_k={tactic.split_k}") + if not _MIN_AB_STAGES <= tactic.ab_stages <= _MAX_AB_STAGES: + raise ValueError( + f"ab_stages must be in [{_MIN_AB_STAGES}, {_MAX_AB_STAGES}], got {tactic.ab_stages}" + ) + if not 1 <= m <= MAX_M: + raise ValueError(f"this low-M policy requires 1 <= M <= {MAX_M}, got {m}") + if n <= 0: + raise ValueError(f"N must be positive, got {n}") + if k <= 0 or k % _CTA_K or (k // _CTA_K) % tactic.split_k: + raise ValueError( + f"K={k} with CTA_K={_CTA_K} does not divide evenly across split_k={tactic.split_k}" + ) + smem_bytes = _smem_bytes(tactic, tactic.ab_stages) + if smem_bytes > smem_capacity: + raise ValueError( + f"tactic {tactic} needs {smem_bytes} B of shared memory but only " + f"{smem_capacity} B are available; max ab_stages is " + f"{_max_ab_stages_for(tactic, smem_capacity)}" + ) + + +@functools.cache +def device_limits(device_index: int) -> tuple[int, int]: + """(SM count, opt-in shared memory per block) of a CUDA device, the resources + ``splitk_default_tactic`` plans for.""" + properties = _torch.cuda.get_device_properties(device_index) + return properties.multi_processor_count, properties.shared_memory_per_block_optin + + +def default_tactic(m: int, n: int, k: int) -> SplitKTactic: + """Choose the default occupancy-oriented tactic.""" + if n <= 512: + mma_m = 64 + mma_n = 16 if n == 512 and m > 24 else 8 + requested_split = 4 + else: + mma_n = 8 if m <= 8 else 16 if m <= 16 else 32 + if n <= 3072: + mma_m = 128 if m <= 16 else 64 + requested_split = 4 if m <= 16 else 2 + elif n < 8192: + mma_m = 64 + requested_split = 2 + else: + mma_m = 128 if k <= 1024 and m <= 24 else 64 + requested_split = 1 + + if k <= 4 * _CTA_K: + requested_split = 1 + split_k = next( + split_k + for split_k in reversed(_SUPPORTED_SPLIT_K) + if split_k <= requested_split and (k // _CTA_K) % split_k == 0 + ) + tactic = SplitKTactic(mma_m, mma_n, split_k, _MIN_AB_STAGES) + max_stages = _max_ab_stages_for(tactic, _SMEM_CAPACITY_BYTES) + tactic = dataclasses.replace( + tactic, + ab_stages=( + _MIN_AB_STAGES + if k == _CTA_K + else _DEFAULT_MIN_AB_STAGES + if k <= 2 * _CTA_K and m > 8 + else min(max_stages, 6) + if k <= 4 * _CTA_K + else max_stages + ), + ) + validate_tactic(tactic, m, n, k) + return tactic + + +def splitk_default_tactic( + m: int, n: int, k: int, *, sm_count: int, smem_capacity: int +) -> SplitKTactic: + """The source default tactic, refined to split 8 where split 8 fills otherwise idle SMs. + + The refinement needs exact K slices of at least two K tiles per peer and at most one wave + of eight-CTA clusters; other shapes keep the source default. An occupancy heuristic, not a + measured winner for every shape. The launch reserves the full SM100 shared-memory capacity, + even when the selected pipeline uses fewer bytes. + """ + if sm_count <= 0 or smem_capacity < _SMEM_CAPACITY_BYTES: + raise ValueError("device cannot serve the split-K launch resources") + tactic = default_tactic(m, n, k) + if tactic.split_k > sm_count: + raise ValueError("device has too few SMs for the split-K cluster") + output_tiles = ((n + tactic.mma_m - 1) // tactic.mma_m) * ( + (m + tactic.mma_n - 1) // tactic.mma_n + ) + # Require exact K slices, at least two tiles per peer, and at most one + # wave of eight-CTA clusters. Wide outputs retain their existing tactic. + if ( + k % (8 * _CTA_K) + or k // (8 * _CTA_K) < _DEFAULT_MIN_AB_STAGES + or output_tiles * tactic.split_k >= sm_count + or output_tiles * 8 > sm_count + ): + return tactic + candidate = dataclasses.replace(tactic, split_k=8) + stages = min( + tactic.ab_stages, + k // (8 * _CTA_K), + _max_ab_stages_for(candidate, _SMEM_CAPACITY_BYTES), + ) + if stages < _DEFAULT_MIN_AB_STAGES: + return tactic + candidate = dataclasses.replace(candidate, ab_stages=stages) + validate_tactic(candidate, m, n, k) + return candidate + + +__all__ = [ + "MAX_M", + "HnormB", + "LamportPublish", + "SplitKTactic", + "default_tactic", + "device_limits", + "hidden_half_norm_supported", + "hnorm_b_supported", + "operands_overlap", + "publish_supported", + "publish_tactic", + "run_splitk_dense", + "run_splitk_hidden_half", + "run_splitk_publish", + "splitk_default_tactic", + "tactic_launch_smem_bytes", + "validate_tactic", +] + + +@dsl_user_op +def _map_shared_rank( + smem_ptr: "cute.Pointer", + peer_cta_rank_in_cluster: "Int32", + *, + loc=None, + ip=None, +) -> "Int32": + """Map an SMEM pointer into a peer CTA's address space.""" + return Int32( + llvm.inline_asm( + T.i32(), + [ + smem_ptr.toint(loc=loc, ip=ip).ir_value(), + peer_cta_rank_in_cluster.ir_value(), + ], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _store_shared_remote_v4( + value0, + value1, + value2, + value3, + smem_ptr: "cute.Pointer", + mbar_ptr: "cute.Pointer", + peer_cta_rank_in_cluster: "Int32", + *, + loc=None, + ip=None, +) -> None: + """Publish four FP32 partials into a peer's SMEM, crediting 16 bytes.""" + llvm.inline_asm( + None, + [ + _map_shared_rank(smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), + value0.bitcast(Int32).ir_value(loc=loc, ip=ip), + value1.bitcast(Int32).ir_value(loc=loc, ip=ip), + value2.bitcast(Int32).ir_value(loc=loc, ip=ip), + value3.bitcast(Int32).ir_value(loc=loc, ip=ip), + _map_shared_rank(mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), + ], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 " + "[$0], {$1, $2, $3, $4}, [$5];", + "r,r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _lamport_current_buffer( + flags_address: "Int64", + buffers_address: "Int64", + *, + loc=None, + ip=None, +) -> "Int64": + """The address of the workspace's current Lamport buffer: buffers + flags[0] (current index) + * flags[2] (bytes per buffer), the flags read with one volatile 16-byte load (the previous + all-reduce rotated them).""" + return Int64( + llvm.inline_asm( + T.i64(), + [flags_address.ir_value(loc=loc, ip=ip), buffers_address.ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 lfcur, lfdirty, lfbytes, lfstages;\n\t.reg .b64 lfoffset;\n\t" + "ld.volatile.global.v4.u32 {lfcur, lfdirty, lfbytes, lfstages}, [$1];\n\t" + "mul.wide.u32 lfoffset, lfcur, lfbytes;\n\tadd.s64 $0, $2, lfoffset;\n\t}", + "=l,l,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _cvt_bf16x2_lamport(lo: "Float32", hi: "Float32", *, loc=None, ip=None) -> "Int32": + """Round two FP32 values to BF16 (nearest even), ``lo`` in the low half, and write a -0.0 half + (0x8000, the idle value of a Lamport buffer) as +0.0, as the one-shot all-reduce does before it + publishes its shard.""" + return Int32( + llvm.inline_asm( + T.i32(), + [lo.ir_value(loc=loc, ip=ip), hi.ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 lpword, lplo, lphi;\n\t.reg .pred lpzlo, lpzhi;\n\t" + "cvt.rn.bf16x2.f32 lpword, $2, $1;\n\t" + "and.b32 lplo, lpword, 0xFFFF;\n\tshr.u32 lphi, lpword, 16;\n\t" + "setp.eq.u32 lpzlo, lplo, 0x8000;\n\tsetp.eq.u32 lpzhi, lphi, 0x8000;\n\t" + "@lpzlo and.b32 lpword, lpword, 0xFFFF0000;\n\t" + "@lpzhi and.b32 lpword, lpword, 0xFFFF;\n\t" + "mov.b32 $0, lpword;\n\t}", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _store_global_v4( + address: "Int64", v0: "Int32", v1: "Int32", v2: "Int32", v3: "Int32", *, loc=None, ip=None +) -> None: + """One 16-byte global store (also through a multicast address).""" + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip)] + [v.ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "st.global.v4.b32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _ld_global_cg_v4(address: "Int64", *, loc=None, ip=None): + """One 16-byte load through L2 (``ld.global.cg``): four words.""" + res = llvm.inline_asm( + llvm.StructType.get_literal([T.i32(), T.i32(), T.i32(), T.i32()]), + [address.ir_value(loc=loc, ip=ip)], + "ld.global.cg.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + return [Int32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(4)] + + +@dsl_user_op +def _st_shared_v4( + address: "Int32", v0: "Int32", v1: "Int32", v2: "Int32", v3: "Int32", *, loc=None, ip=None +) -> None: + """One 16-byte shared-memory store.""" + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip)] + [v.ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "st.shared.v4.b32 [$0], {$1, $2, $3, $4};", + "r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _fence_proxy_async_shared(*, loc=None, ip=None) -> None: + """Order this thread's shared-memory stores before later async-proxy (tensor core) reads.""" + llvm.inline_asm( + None, + [], + "fence.proxy.async.shared::cta;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _add_bf16x2_squares(acc: "Float32", word: "Int32", *, loc=None, ip=None) -> "Float32": + """``(acc + lo * lo) + hi * hi`` for the two BF16 halves of ``word`` (low half first), every + product and sum rounded to nearest, no FMA: FlashInfer RMSNorm's per-thread accumulation.""" + return Float32( + llvm.inline_asm( + T.f32(), + [acc.ir_value(loc=loc, ip=ip), word.ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 sqlo, sqhi;\n\t.reg .f32 sqa, sqb, sqp, sqs;\n\t" + "shl.b32 sqlo, $2, 16;\n\tand.b32 sqhi, $2, 0xFFFF0000;\n\t" + "mov.b32 sqa, sqlo;\n\tmov.b32 sqb, sqhi;\n\t" + "mul.rn.f32 sqp, sqa, sqa;\n\tadd.rn.f32 sqs, $1, sqp;\n\t" + "mul.rn.f32 sqp, sqb, sqb;\n\tadd.rn.f32 $0, sqs, sqp;\n\t}", + "=f,f,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _add_rn(a: "Float32", b: "Float32", *, loc=None, ip=None) -> "Float32": + """``a + b`` rounded to nearest (never contracted).""" + return Float32( + llvm.inline_asm( + T.f32(), + [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +def _f32_hex(value: float) -> str: + """PTX literal of the FP32 rounding of ``value``.""" + return "0f%08X" % struct.unpack(" "Float32": + """FlashInfer RMSNorm's ``rsqrt.approx.ftz(total / hidden + eps)``, the division and the sum + rounded to nearest (IEEE).""" + return Float32( + llvm.inline_asm( + T.f32(), + [total.ir_value(loc=loc, ip=ip)], + "{\n\t.reg .f32 rsm, rst;\n\t" + f"div.rn.f32 rsm, $1, {_f32_hex(float(hidden))};\n\t" + f"add.rn.f32 rst, rsm, {_f32_hex(eps)};\n\t" + "rsqrt.approx.ftz.f32 $0, rst;\n\t}", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _rmsnorm_bf16x2( + x_word: "Int32", w_word: "Int32", rstd: "Float32", *, loc=None, ip=None +) -> "Int32": + """FlashInfer RMSNorm's output for two BF16 elements (low half first): ``(x * rstd) * (w + 0)`` + rounded to nearest at every step, then to BF16 (nearest even).""" + return Int32( + llvm.inline_asm( + T.i32(), + [ + x_word.ir_value(loc=loc, ip=ip), + w_word.ir_value(loc=loc, ip=ip), + rstd.ir_value(loc=loc, ip=ip), + ], + "{\n\t.reg .b32 rnxl, rnxh, rnwl, rnwh;\n\t" + ".reg .f32 rfxl, rfxh, rfwl, rfwh, ryl, ryh;\n\t" + "shl.b32 rnxl, $1, 16;\n\tand.b32 rnxh, $1, 0xFFFF0000;\n\t" + "shl.b32 rnwl, $2, 16;\n\tand.b32 rnwh, $2, 0xFFFF0000;\n\t" + "mov.b32 rfxl, rnxl;\n\tmov.b32 rfxh, rnxh;\n\t" + "mov.b32 rfwl, rnwl;\n\tmov.b32 rfwh, rnwh;\n\t" + "add.rn.f32 rfwl, rfwl, 0f00000000;\n\tadd.rn.f32 rfwh, rfwh, 0f00000000;\n\t" + "mul.rn.f32 ryl, rfxl, $3;\n\tmul.rn.f32 ryl, ryl, rfwl;\n\t" + "mul.rn.f32 ryh, rfxh, $3;\n\tmul.rn.f32 ryh, ryh, rfwh;\n\t" + "cvt.rn.bf16x2.f32 $0, ryh, ryl;\n\t}", + "=r,r,r,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +#: Rank that gathers partials and stores the output. +OWNER_RANK = 0 + + +class SplitKDenseGemmKernel: + """Standalone BF16/FP16 GEMM with a cluster-local split-K reduction.""" + + def __init__( + self, + *, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + reserve_full_smem: bool = True, + publish: LamportPublish | None = None, + hnorm_b: HnormB | None = None, + ) -> None: + self.acc_dtype = cutlass.Float32 + self.cta_m = tactic.mma_m + self.cta_n = tactic.mma_n + self.cta_k = _CTA_K + self.num_ab_stage = tactic.ab_stages + self.split_k = tactic.split_k + self.use_pdl = use_pdl + self.has_bias = has_bias + # Shared memory requested at launch: the full capacity (one CTA per SM) or only the + # tactic's footprint. The kernel itself is the same either way. + self.reserve_full_smem = reserve_full_smem + self.launch_smem_bytes = tactic_launch_smem_bytes(tactic) + # Lamport publish (run_splitk_publish): the owner stages the reduced FP32 tile in shared + # memory and stores it into the all-reduce workspace instead of writing C. + self.publish = publish + # HnormB: warp 1 computes B (FlashInfer's RMSNorm of the activation row) instead of + # loading it; the norm weight takes the bias operand's place. + self.hnorm_b = hnorm_b + if hnorm_b is not None and not hnorm_b_supported(tactic, hnorm_b.hidden): + raise ValueError( + f"HnormB needs hidden {HNORM_B_HIDDEN} and a tactic of hnorm_b_supported: " + f"{tactic}, hidden {hnorm_b.hidden}" + ) + if hnorm_b is not None and (has_bias or publish is not None): + raise ValueError("HnormB excludes a bias and the Lamport publish") + if publish is not None: + if has_bias or self.cta_m % _PUBLISH_VEC_ELEMENTS: + raise ValueError(f"Lamport publish needs no bias and 8-column tiles: {tactic}") + self.launch_smem_bytes = min( + _align_up(_publish_smem_bytes(tactic, tactic.ab_stages), 1024) + + _LAUNCH_SMEM_HEADROOM_BYTES, + _SMEM_CAPACITY_BYTES, + ) + + self.threads_per_cta = 256 + self.epilog_threads = 128 + self.mma_tiler_mn = (tactic.mma_m, tactic.mma_n) + self.cta_group = tcgen05.CtaGroup.ONE + self.tma_op = cute_ext.OperationTypeEnum.SM90_TMA_LOAD + self.cluster_shape = (1, tactic.split_k, 1) + + values_per_thread = (tactic.mma_m * tactic.mma_n) // self.epilog_threads + if values_per_thread % 4: + raise ValueError( + f"CTA tile ({tactic.mma_m}, {tactic.mma_n}) gives " + f"{values_per_thread} " + "values per epilogue thread; remote stores require a multiple of 4" + ) + self.mailbox_elements = (tactic.split_k - 1) * self.epilog_threads * values_per_thread + self.expected_transaction_bytes = self.mailbox_elements * _FP32_BYTES + + @cute.experimental.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + lamport: cute.Tensor, + flags: cute.Tensor, + stream: _cuda.CUstream, + ): + # Grid-y packs output-N tile and cluster rank. + self.kernel(a, b, c, bias, lamport, flags).launch( + grid=( + cute.ceil_div(c.layout.shape[0], self.cta_m), + cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k, + c.layout.shape[2], + ), + block=(self.threads_per_cta, 1, 1), + cluster=self.cluster_shape, + smem=cute.Int64( + utils.get_smem_capacity_in_bytes("sm_100") + if self.reserve_full_smem + else self.launch_smem_bytes + ), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.experimental.kernel + def kernel( + self, + mA: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L), K-major + mB: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L), K-major + mC: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L), M-major + mBias: cute.Tensor, # Broadcast bias; dead when has_bias=False + mLamport: cute.Tensor, # BF16 all-reduce workspace, multicast address; dead without publish + mFlags: cute.Tensor, # The workspace's buffer flags; dead without publish + ): + """Allocate storage and dispatch the specialized warps.""" + stages = self.num_ab_stage + + ab_dtype = mA.element_type + tiled_mma = sm100_utils.make_trivial_tiled_mma( + ab_dtype, + ab_dtype, + utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + + mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k) + block_idx = cute.arch.block_idx() + bidx = block_idx[0] + split_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + n_idx = block_idx[1] // self.split_k + l_idx = block_idx[2] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + sA = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + sB = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + + acc_layout = cute_ext.make_tmem_layout_acc(tiled_mma, self.mma_tiler_mn, acc_stage=1) + c_tiler_mn = (self.cta_m, self.cta_n) + + bar_full = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_empty = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_mma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tmem_alloc = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + tmem_base_ptr = cute_ext.allocate( + cutlass.Int32, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_TMEM_POINTER_BYTES, + ).iterator + + if cutlass.const_expr(self.split_k > 1): + mailbox = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.mailbox_elements), + alignment=_MAILBOX_ALIGN_BYTES, + ) + bar_reduce = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + else: + # Dummy operands for the compile-time-elided reduction. + mailbox = sA + bar_reduce = bar_mma_epilog + + if cutlass.const_expr(self.publish is not None): + # Lamport publish: the reduced FP32 tile, M-major as the output tile. + sPub = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout((self.cta_m, self.cta_n)), + alignment=_PUBLISH_ALIGN_BYTES, + ) + else: + # Dummy operand for the compile-time-elided publish. + sPub = sA + + if warp_idx == 0: + with cute.arch.elect_one(): + for i in range(stages): + cute.arch.mbarrier_init(bar_full + i, 2) + cute.arch.mbarrier_init(bar_empty + i, 1) + cute.arch.mbarrier_init(bar_tma_epilog, 32) + cute.arch.mbarrier_init(bar_mma_epilog, 1) + cute.arch.mbarrier_init(bar_tmem_alloc, 160) + + if cutlass.const_expr(self.split_k > 1): + # Owner arrival plus peer transaction-byte credits. + cute.arch.mbarrier_init(bar_reduce, 1) + + cute.arch.mbarrier_init_fence() + if cutlass.const_expr(self.split_k > 1): + # Publish peer barriers before cross-CTA stores. + cute.arch.cluster_arrive_relaxed() + else: + cute.arch.barrier() + + # Host validation guarantees an equal, tail-free K partition. + k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k + k_tile_start = split_rank * k_tile_count + + if cutlass.const_expr(self.split_k > 1): + cute.arch.cluster_wait() + + # Warp 3 is idle; warps 4-7 run the epilogue. + if warp_idx == 0: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), + sA, + cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), + k_tile_start, + k_tile_count, + True, + ) + elif warp_idx == 1: + if cutlass.const_expr(self.hnorm_b is not None): + self.hnorm_b_warp(bar_full, bar_empty, sB, mB, mBias, k_tile_start) + else: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), + sB, + cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), + k_tile_start, + k_tile_count, + False, + ) + elif warp_idx == 2: + self.mma_warp( + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma, + sA, + sB, + tmem_base_ptr, + acc_layout, + self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), + k_tile_count, + ) + elif warp_idx >= 4: + self.epilog_warp( + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout, + cute.local_tile(mC, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.local_tile(mBias, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.arch.thread_idx()[0] - 128, + mC.element_type, + utils.LayoutEnum.from_tensor(mC), + mailbox, + bar_reduce, + split_rank, + sPub, + mC, + mLamport, + mFlags, + bidx, + n_idx, + ) + + @cute.experimental.jit + def dma_warp( + self, + bar_full, + bar_empty, + bar_tma_epilog, + g_tile: cute.Tensor, + s_tile: cute.Tensor, + cta_v_map: cute.Layout, + k_tile_start: cutlass.Int32, + k_tile_count: cutlass.Int32, + is_a: cutlass.Constexpr, + ): + stages = self.num_ab_stage + if cutlass.const_expr(not is_a and self.use_pdl): + cute.arch.griddepcontrol_wait() + + empty_phase = cutlass.Int32(1) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + bar_full + stage, + cute.size_in_bytes( + s_tile.element_type, + cute.slice_(s_tile.layout, (None, None, None, 0)), + ), + ) + cute_ext.tma_load( + g_tile[None, None, k_tile_start + k_tile], + s_tile[None, None, None, stage], + (bar_full + stage).value, + cta_v_map=cta_v_map, + tma_operation_type=self.tma_op, + update_expect_tx=False, + ) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + if cutlass.const_expr( + is_a + and self.use_pdl + and not (self.publish is not None and self.publish.trigger_after_wait) + ): + cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(not is_a and self.has_bias): + cute.arch.mbarrier_arrive(bar_tma_epilog) + self._drain_producer(bar_empty, empty_phase, k_tile_count) + + @cute.experimental.jit + def hnorm_b_warp( + self, + bar_full, + bar_empty, + sB: cute.Tensor, + mX: cute.Tensor, + mW: cute.Tensor, + k_tile_start: cutlass.Int32, + ): + """B of HnormB: FlashInfer's RMSNorm of the activation row x with weight w, written into + this CTA's B stages where the TMA would have put the norm's output. + + FlashInfer's CuTe kernel (H = 6144) gives a row to 64 threads; thread t holds the 8-element + vectors t + 64 b (b = 0..11), sums their squares in order, and the row sum is the butterfly + sum of threads 0-31 plus that of threads 32-63. Lane l of this warp plays threads l and + l + 32: it holds the vectors l + 32 m (m = 0..23; even m are thread l's, odd m thread + l + 32's). The B tiles are 8 x 128 K-major with the K_SW128 layout: row 0 (the activation; + rows 1-7 are zero, as the TMA fills rows past M) is unswizzled, K columns [0, 64) at bytes + [0, 128) and [64, 128) at bytes [1024, 1152) of the stage.""" + hidden = self.hnorm_b.hidden + lane = cute.arch.lane_idx() + vec_per_tile = self.cta_k // _HNORM_B_VEC_ELEMENTS + k_tiles = hidden // self.cta_k // self.split_k + per_lane = k_tiles * vec_per_tile // _WARP_LANES + row_rounds = hidden // _HNORM_B_VEC_ELEMENTS // _WARP_LANES + stage_bytes = self.cta_n * self.cta_k * _AB_ELEMENT_BYTES + atom_bytes = stage_bytes // 2 + row_bytes = atom_bytes // self.cta_n + sb = sB.iterator.toint() + zero = Int32(0) + + # Rows 1-7 of the stages this CTA uses. + zero_vecs = (atom_bytes - row_bytes) // 16 + for s in cutlass.range_constexpr(k_tiles): + for r in cutlass.range_constexpr((2 * zero_vecs + _WARP_LANES - 1) // _WARP_LANES): + idx = lane + Int32(r * _WARP_LANES) + if idx < 2 * zero_vecs: + atom = idx // zero_vecs + _st_shared_v4( + sb + + Int32(s * stage_bytes + row_bytes) + + atom * atom_bytes + + (idx - atom * zero_vecs) * 16, + zero, + zero, + zero, + zero, + ) + + # This CTA's K slice is the vectors [v0, v0 + k_tiles * vec_per_tile); lane l produces the + # slice vectors v0 + l + 32 j. Their norm weights are constants: loaded before the wait. + v0 = k_tile_start * vec_per_tile + w_words = [] + for j in cutlass.range_constexpr(per_lane): + v = v0 + lane + Int32(j * _WARP_LANES) + w_words.append(_ld_global_cg_v4((mW.iterator + v * _HNORM_B_VEC_ELEMENTS).toint())) + + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + + x_words = [] + for m in cutlass.range_constexpr(row_rounds): + v = lane + Int32(m * _WARP_LANES) + x_words.append(_ld_global_cg_v4((mX.iterator + v * _HNORM_B_VEC_ELEMENTS).toint())) + xs_words = [] + for j in cutlass.range_constexpr(per_lane): + v = v0 + lane + Int32(j * _WARP_LANES) + xs_words.append(_ld_global_cg_v4((mX.iterator + v * _HNORM_B_VEC_ELEMENTS).toint())) + + acc_lo = Float32(0.0) + acc_hi = Float32(0.0) + for b in cutlass.range_constexpr(row_rounds // 2): + for q in cutlass.range_constexpr(4): + acc_lo = _add_bf16x2_squares(acc_lo, x_words[2 * b][q]) + for b in cutlass.range_constexpr(row_rounds // 2): + for q in cutlass.range_constexpr(4): + acc_hi = _add_bf16x2_squares(acc_hi, x_words[2 * b + 1][q]) + for offset in (1, 2, 4, 8, 16): + acc_lo = _add_rn(acc_lo, cute.arch.shuffle_sync_bfly(acc_lo, offset)) + acc_hi = _add_rn(acc_hi, cute.arch.shuffle_sync_bfly(acc_hi, offset)) + rstd = _rmsnorm_rstd(_add_rn(acc_lo, acc_hi), hidden, self.hnorm_b.eps) + + for j in cutlass.range_constexpr(per_lane): + local = lane + Int32(j * _WARP_LANES) + tile = local // vec_per_tile + col_vec = local - tile * vec_per_tile + col_atom = col_vec // (vec_per_tile // 2) + _st_shared_v4( + sb + + tile * stage_bytes + + col_atom * atom_bytes + + (col_vec - col_atom * (vec_per_tile // 2)) * 16, + _rmsnorm_bf16x2(xs_words[j][0], w_words[j][0], rstd), + _rmsnorm_bf16x2(xs_words[j][1], w_words[j][1], rstd), + _rmsnorm_bf16x2(xs_words[j][2], w_words[j][2], rstd), + _rmsnorm_bf16x2(xs_words[j][3], w_words[j][3], rstd), + ) + # The tensor core reads the stages through the async proxy. + _fence_proxy_async_shared() + cute.arch.sync_warp() + with cute.arch.elect_one(): + for s in cutlass.range_constexpr(k_tiles): + cute.arch.mbarrier_arrive(bar_full + s) + # As the TMA producer: wait until the MMA has released every stage this CTA filled. + self._drain_producer(bar_empty, Int32(1 ^ (k_tiles // self.num_ab_stage)), Int32(k_tiles)) + + @cute.experimental.jit + def _drain_producer( + self, + bar_empty, + empty_phase: cutlass.Int32, + k_tile_count: cutlass.Int32, + ): + stages = self.num_ab_stage + for tail in cutlass.range(stages, unroll=1): + stage = (tail + k_tile_count) % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + @cute.experimental.jit + def mma_warp( + self, + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma: cute.TiledMma, + sA: cute.Tensor, + sB: cute.Tensor, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + mma_inst_tile_k: cutlass.Constexpr, + k_tile_count: cutlass.Int32, + ): + num_tmem_cols = 256 + cute.arch.alloc_tmem(num_tmem_cols, tmem_base_ptr, is_two_cta=False) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False) + + tmem_ptr = cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr) + accumulator = cute.make_tensor(tmem_ptr, acc_layout)[None, None, None, 0] + mma_atom = cute.make_mma_atom(tiled_mma.op) + full_phase = cutlass.Int32(0) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % self.num_ab_stage + cute.arch.mbarrier_wait(bar_full + stage, full_phase) + for k_block in range(mma_inst_tile_k): + if k_block == 0: + mma_atom.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + else: + mma_atom.set(tcgen05.Field.ACCUMULATE, True) + cute_ext.dot( + mma_atom, + cute.append_ones(sA[None, None, k_block, stage], up_to_rank=3), + cute.append_ones(sB[None, None, k_block, stage], up_to_rank=3), + accumulator, + ) + with cute.arch.elect_one(): + tcgen05.commit(bar_empty + stage, None, self.cta_group) + if stage == self.num_ab_stage - 1: + full_phase = full_phase ^ 1 + + with cute.arch.elect_one(): + tcgen05.commit(bar_mma_epilog, None, self.cta_group) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 1) + cute.arch.dealloc_tmem(tmem_ptr, num_tmem_cols, is_two_cta=False) + + @cute.experimental.jit + def epilog_warp( + self, + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + gD_tile: cute.Tensor, + gBias_tile: cute.Tensor, + epi_tid: cutlass.Int32, + c_dtype: cutlass.Constexpr, + d_layout: cutlass.Constexpr, + mailbox, + bar_reduce, + split_rank: cutlass.Int32, + sPub, + mC: cute.Tensor, + mLamport: cute.Tensor, + mFlags: cute.Tensor, + bidx: cutlass.Int32, + n_idx: cutlass.Int32, + ): + lamport_buffer = Int64(0) + if cutlass.const_expr(self.publish is not None): + # The flags are read after the grid dependency wait: the previous all-reduce, which + # rotated them, precedes this grid's primary. Every CTA's epilogue warps wait; with + # trigger_after_wait they also release the dependents here. + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(self.publish.trigger_after_wait): + cute.arch.griddepcontrol_launch_dependents() + lamport_buffer = _lamport_current_buffer( + mFlags.iterator.toint(), mLamport.iterator.toint() + ) + + # Wait until MMA publishes the TMEM base pointer. + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) + + acc_view = cute.make_tensor( + cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr), + acc_layout, + )[((None, None), 0, 0, 0)] + + epi_tile = (self.cta_m, self.cta_n) + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( + sm100_utils.get_tmem_load_op( + (self.cta_m, self.cta_n, self.cta_k), + d_layout, + c_dtype, + self.acc_dtype, + epi_tile, + False, + ), + acc_view, + ) + gD_epi = cute.flat_divide(gD_tile, epi_tile) + + # Match each epilogue thread's TMEM partition in RMEM. + rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid) + rAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rD = cute_ext.allocate( + c_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + thr_t2r = tiled_copy_t2r.get_slice(epi_tid) + + if cutlass.const_expr(self.has_bias): + bias_dtype = gBias_tile.element_type + rBias = cute_ext.allocate( + bias_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rBiasAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + if split_rank == OWNER_RANK: + cute.arch.mbarrier_wait(bar_tma_epilog, 0) + cute_ext.partition_and_copy( + cute.make_tiled_copy_D( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), bias_dtype), + tiled_copy_t2r, + ).get_slice(epi_tid), + cute.flat_divide(gBias_tile, epi_tile)[None, None, 0, 0], + rBias, + ) + rBiasAcc.store(rBias.load().to(self.acc_dtype)) + + cute.arch.mbarrier_wait(bar_mma_epilog, 0) + cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) + # Make tcgen05.ld visible before TMEM release and RMEM use. + cute.arch.fence_view_async_tmem_load() + cute.arch.mbarrier_arrive(bar_tmem_alloc) + + # Peers publish FP32 partials; only rank 0 reduces and stores. + if cutlass.const_expr(self.split_k > 1): + assert cute.size(rmem_layout) == self.mailbox_elements // ( + (self.split_k - 1) * self.epilog_threads + ) + values_per_thread = cutlass.const_expr(cute.size(rmem_layout)) + values_per_peer = cutlass.const_expr(self.epilog_threads * values_per_thread) + if split_rank != OWNER_RANK: + for value_idx in cutlass.range_constexpr(0, values_per_thread, 4): + _store_shared_remote_v4( + rAcc[value_idx], + rAcc[value_idx + 1], + rAcc[value_idx + 2], + rAcc[value_idx + 3], + mailbox.iterator + + (split_rank - Int32(1)) * values_per_peer + + epi_tid * values_per_thread + + value_idx, + bar_reduce, + Int32(OWNER_RANK), + ) + else: + if epi_tid == 0: + cute.arch.mbarrier_arrive_and_expect_tx( + bar_reduce, self.expected_transaction_bytes + ) + cute.arch.mbarrier_wait(bar_reduce, 0) + for peer in cutlass.range_constexpr(self.split_k - 1): + for value_idx in cutlass.range_constexpr(values_per_thread): + rAcc[value_idx] = ( + rAcc[value_idx] + + mailbox[ + peer * values_per_peer + epi_tid * values_per_thread + value_idx + ] + ) + + if split_rank == OWNER_RANK: + if cutlass.const_expr(self.publish is not None): + # Stage the tile by output coordinate, then store eight columns per thread. + cute_ext.partition_and_copy(thr_t2r, rAcc, sPub) + cute.arch.barrier( + barrier_id=_PUBLISH_BARRIER_ID, number_of_threads=self.epilog_threads + ) + self._publish_tile(sPub, mC, lamport_buffer, bidx, n_idx, epi_tid) + else: + if cutlass.const_expr(self.has_bias): + rAcc.store(rAcc.load() + rBiasAcc.load()) + + rD.store(rAcc.load().to(c_dtype)) + # Preserve TMEM coordinates; the copy predicates output tails. + cute_ext.partition_and_copy(thr_t2r, rD, gD_epi[None, None, 0, 0]) + + # The reduction mbarrier covers remote stores; no cluster barrier needed. + + @cute.experimental.jit + def _publish_tile( + self, + sPub: cute.Tensor, + mC: cute.Tensor, + lamport_buffer: Int64, + bidx: cutlass.Int32, + n_idx: cutlass.Int32, + epi_tid: cutlass.Int32, + ): + """Store the staged tile into the Lamport buffer: eight consecutive output columns of one + row per thread, rounded to BF16 as C would be, -0.0 as +0.0, one 16-byte store each at + [token][rank][column] (the one-shot all-reduce's slot of this rank).""" + vecs_m = cutlass.const_expr(self.cta_m // _PUBLISH_VEC_ELEMENTS) + num_vecs = cutlass.const_expr(vecs_m * self.cta_n) + n_feat = cute.size(mC, mode=[0]) + n_tok = cute.size(mC, mode=[1]) + for it in cutlass.range_constexpr( + (num_vecs + self.epilog_threads - 1) // self.epilog_threads + ): + vec = epi_tid + it * self.epilog_threads + row = vec // vecs_m + m8 = (vec - row * vecs_m) * _PUBLISH_VEC_ELEMENTS + token = n_idx * self.cta_n + row + feature = bidx * self.cta_m + m8 + if vec < num_vecs: + if token < n_tok: + if feature < n_feat: + words = [] + for j in cutlass.range_constexpr(_PUBLISH_VEC_ELEMENTS // 2): + words.append( + _cvt_bf16x2_lamport( + sPub[m8 + 2 * j, row], sPub[m8 + 2 * j + 1, row] + ) + ) + slot = Int64(token * self.publish.world_size + self.publish.rank) + offset = (slot * Int64(n_feat) + Int64(feature)) * 2 + _store_global_v4( + lamport_buffer + offset, words[0], words[1], words[2], words[3] + ) + + +_SUPPORTED_TORCH_DTYPES = (_torch.bfloat16, _torch.float16) + + +@cute.experimental.jit +def _bmm_no_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + stream: _cuda.CUstream, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + # C also stands in for the dead bias and Lamport operands. + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + c, + c, + stream, + ) + + +@cute.experimental.jit +def _bmm_hnorm_b( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + norm_weight: cute.Tensor, + stream: _cuda.CUstream, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + # HnormB: b is the raw activation and the norm weight takes the dead bias operand's place (with + # C's layout: only its address is read); C stands in for the dead Lamport operands. + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(norm_weight.iterator, c.layout), + c, + c, + stream, + ) + + +@cute.experimental.jit +def _bmm_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + stream: _cuda.CUstream, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + # C also stands in for the dead Lamport operands. + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(bias.iterator, cute.select(bias.layout, mode=[1, 2, 0])), + c, + c, + stream, + ) + + +@cute.experimental.jit +def _bmm_publish( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + lamport: cute.Tensor, + flags: cute.Tensor, + stream: _cuda.CUstream, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + # C (not written) also stands in for the dead bias. + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + c, + lamport, + flags, + stream, + ) + + +def _from_dlpack_dynamic(tensor, leading_dim: int, assumed_align: int = 32): + return from_dlpack(tensor, assumed_align=assumed_align).mark_layout_dynamic( + leading_dim=leading_dim + ) + + +def _detect_leading_dim(tensor: _torch.Tensor) -> int: + # Ignore synthetic batch stride, including 1x1. + for dim, stride in enumerate(tensor.stride()[1:], start=1): + if stride == 1: + return dim + raise ValueError("tensor has no stride-1 dimension") + + +def _make_layout_tensor( + shape: tuple[int, ...], dtype: _torch.dtype, leading_dim: int +) -> _torch.Tensor: + permutation = [dim for dim in range(len(shape)) if dim != leading_dim] + [leading_dim] + return _torch.empty( + tuple(shape[dim] for dim in permutation), dtype=dtype, device="cuda" + ).permute([permutation.index(dim) for dim in range(len(shape))]) + + +def _make_compile_repr_tensors( + dtype: _torch.dtype, + has_bias: bool, + a_leading: int, + b_leading: int, + c_leading: int, +): + m, n, k, batch = 64, 8, _CTA_K, 1 + tensors = tuple( + _from_dlpack_dynamic(_make_layout_tensor(shape, dtype, leading_dim), leading_dim) + for shape, leading_dim in ( + ((batch, n, k), a_leading), + ((batch, k, m), b_leading), + ((batch, n, m), c_leading), + ) + ) + if not has_bias: + return (*tensors, None) + return ( + *tensors, + _from_dlpack_dynamic( + _torch.empty((n,), dtype=dtype, device="cuda").as_strided( + size=(batch, n, m), stride=(0, 1, 0) + ), + 1, + 2, + ), + ) + + +def _to_cute_swap(a, b, out, bias): + a_swap = b.unsqueeze(0).transpose(-2, -1) + b_swap = a.unsqueeze(0).transpose(-2, -1) + c_swap = out.unsqueeze(0).transpose(-2, -1) + leading_dims = tuple(_detect_leading_dim(tensor) for tensor in (a_swap, b_swap, c_swap)) + cute_tensors = tuple( + _from_dlpack_dynamic(tensor, leading_dim) + for tensor, leading_dim in zip((a_swap, b_swap, c_swap), leading_dims, strict=True) + ) + if bias is None: + return (*cute_tensors, None, leading_dims) + return ( + *cute_tensors, + _from_dlpack_dynamic( + bias.as_strided(size=(1, c_swap.shape[1], c_swap.shape[2]), stride=(0, 1, 0)), + 1, + 2, + ), + leading_dims, + ) + + +@functools.cache +def _get_compiled_splitk_kernel( + dtype, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + leading_dims: tuple[int, int, int], + reserve_full_smem: bool = True, + publish: LamportPublish | None = None, +): + if dtype not in _SUPPORTED_TORCH_DTYPES: + raise ValueError(f"split-K dense GEMM supports {_SUPPORTED_TORCH_DTYPES}; got {dtype}") + + kernel = SplitKDenseGemmKernel( + tactic=tactic, + use_pdl=use_pdl, + has_bias=has_bias, + reserve_full_smem=reserve_full_smem, + publish=publish, + ) + compile_tensors = _make_compile_repr_tensors(dtype, has_bias, *leading_dims) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + if publish is not None: + lamport = _from_dlpack_dynamic(_torch.empty((64,), dtype=dtype, device="cuda"), 0) + flags = _from_dlpack_dynamic(_torch.empty((16,), dtype=_torch.int32, device="cuda"), 0) + compiled = cute_ext.compile( + _bmm_publish, kernel, *compile_tensors[:3], lamport, flags, stream + ) + elif has_bias: + compiled = cute_ext.compile(_bmm_bias, kernel, *compile_tensors, stream) + else: + compiled = cute_ext.compile(_bmm_no_bias, kernel, *compile_tensors[:3], stream) + return compiled + + +def operands_overlap(a: _torch.Tensor, b: _torch.Tensor) -> bool: + """Whether the byte ranges that ``a`` and ``b`` span on one device intersect.""" + if a.numel() == 0 or b.numel() == 0 or a.device != b.device: + return False + + def span(t: _torch.Tensor) -> int: + extent = sum((size - 1) * abs(stride) for size, stride in zip(t.shape, t.stride())) + return (extent + 1) * t.element_size() + + a_lo, b_lo = a.data_ptr(), b.data_ptr() + return a_lo < b_lo + span(b) and b_lo < a_lo + span(a) + + +def _validate_runtime_tensors(a, b, bias, out) -> tuple[int, int, int]: + tensors = (a, b, out) + ((bias,) if bias is not None else ()) + if any(not isinstance(tensor, _torch.Tensor) for tensor in tensors): + raise ValueError("a, b, out, and bias must be torch tensors") + if a.ndim != 2 or b.ndim != 2 or out.ndim != 2: + raise ValueError("split-K dense GEMM accepts only 2D tensors") + if a.device.type != "cuda" or any(tensor.device != a.device for tensor in tensors): + raise ValueError("all tensors must be on the same CUDA device") + if a.dtype not in _SUPPORTED_TORCH_DTYPES or any(tensor.dtype != a.dtype for tensor in tensors): + raise ValueError("a, b, out, and bias must share BF16 or FP16 dtype") + if any(not (tensor.is_contiguous() or tensor.t().is_contiguous()) for tensor in (a, b, out)): + raise ValueError("a, b, and out must be dense row-major or column-major matrices") + if any(tensor.data_ptr() % 32 for tensor in (a, b, out)): + raise ValueError("a, b, and out must be 32-byte aligned") + # cute_ext.tma_load is used for a and b (not out, which uses partition_and_copy). + # TMA requires the non-innermost byte stride to be a multiple of 16 bytes. + # For a dense 2D tensor this is max(stride(0), stride(1)) * element_size. + _TMA_STRIDE_ALIGN = 16 + for tensor, tname in ((a, "a"), (b, "b")): + outer_stride_bytes = max(tensor.stride(0), tensor.stride(1)) * tensor.element_size() + if outer_stride_bytes % _TMA_STRIDE_ALIGN: + raise ValueError( + f"tensor '{tname}' outer byte stride {outer_stride_bytes} is not a " + f"multiple of {_TMA_STRIDE_ALIGN}; TMA load requires " + f"{_TMA_STRIDE_ALIGN}-byte-aligned strides " + f"(shape={tuple(tensor.shape)}, strides={tensor.stride()})." + ) + + m, k = a.shape + if b.shape[0] != k: + raise ValueError(f"incompatible shapes: a is {tuple(a.shape)}, b is {tuple(b.shape)}") + n = b.shape[1] + if out.shape != (m, n): + raise ValueError(f"out must have shape {(m, n)}, got {tuple(out.shape)}") + if bias is not None and (bias.ndim != 1 or bias.shape[0] != n or not bias.is_contiguous()): + raise ValueError( + f"bias must be contiguous with shape {(n,)}, " + f"got shape {tuple(bias.shape)} and stride {bias.stride()}" + ) + + return m, n, k + + +def run_splitk_dense( + a: _torch.Tensor, + b: _torch.Tensor, + bias: _torch.Tensor | None, + out: _torch.Tensor, + pdl: bool, + tactic: SplitKTactic, + reserve_full_smem: bool = True, +) -> _torch.Tensor: + """Run ``A[M,K] @ B[K,N]`` with the split-K BF16/FP16 GEMM kernel. + + Args: + a: Input tensor of shape ``[M, K]``, row-major BF16/FP16. + b: Weight tensor of shape ``[K, N]``, column-major BF16/FP16 + (i.e. the transpose of a ``[N, K]`` contiguous weight matrix). + bias: Optional bias of shape ``[N]``, or ``None``. + out: Pre-allocated output tensor of shape ``[M, N]``, BF16/FP16. + pdl: Whether to use Programmatic Dependent Launch. + tactic: The ``SplitKTactic`` specifying tile and pipeline parameters. + reserve_full_smem: Reserve the full shared-memory capacity per CTA (one CTA per + SM). ``False`` reserves only ``tactic_launch_smem_bytes(tactic)``; same kernel + and result. + + Returns: + The ``out`` tensor, filled with the GEMM result. + """ + validate_tactic(tactic, *_validate_runtime_tensors(a, b, bias, out)) + has_bias = bias is not None + cute_tensors = _to_cute_swap(a, b, out, bias) + compiled = _get_compiled_splitk_kernel( + dtype=a.dtype, + tactic=tactic, + use_pdl=pdl, + has_bias=has_bias, + leading_dims=cute_tensors[4], + reserve_full_smem=reserve_full_smem, + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + if has_bias: + compiled(*cute_tensors[:4], stream) + else: + compiled(*cute_tensors[:3], stream) + return out + + +def publish_tactic( + tactic: SplitKTactic, *, smem_capacity: int = _SMEM_CAPACITY_BYTES +) -> SplitKTactic: + """``tactic`` with at most as many A/B stages as fit beside the publish staging tile (the + same results; only the pipeline depth changes).""" + stages = tactic.ab_stages + while stages > 1 and _publish_smem_bytes(tactic, stages) > smem_capacity: + stages -= 1 + return dataclasses.replace(tactic, ab_stages=stages) + + +def publish_supported( + m: int, + n: int, + tactic: SplitKTactic, + *, + dtype: _torch.dtype = _torch.bfloat16, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> bool: + """Whether ``run_splitk_publish`` serves an ``[m, n]`` BF16 output with ``tactic``: 16-byte + column groups and the staging tile beside the A/B stages.""" + return ( + dtype == _torch.bfloat16 + and 1 <= m <= MAX_M + and n > 0 + and n % _PUBLISH_VEC_ELEMENTS == 0 + and tactic.mma_m % _PUBLISH_VEC_ELEMENTS == 0 + and _publish_smem_bytes(tactic, tactic.ab_stages) <= smem_capacity + ) + + +def run_splitk_publish( + a: _torch.Tensor, + b: _torch.Tensor, + out: _torch.Tensor, + lamport: _torch.Tensor, + flags: _torch.Tensor, + pdl: bool, + tactic: SplitKTactic, + publish: LamportPublish, + reserve_full_smem: bool = True, +) -> None: + """Run ``A[M,K] @ B[K,N]`` (no bias) and publish the BF16 result into an MNNVL one-shot + all-reduce workspace instead of writing ``out``. + + Args: + a, b, tactic, pdl, reserve_full_smem: as for ``run_splitk_dense``. + out: The ``[M, N]`` output ``run_splitk_dense`` would write; it only gives the output + layout and is not written. + lamport: The workspace's multicast view (BF16, 1-D, its NUM_LAMPORT_BUFFERS buffers). + flags: The workspace's buffer flags viewed as int32. + publish: World size and rank of the all-reduce group, and the trigger placement. + + Every value lands at ``[token][publish.rank][column]`` of the buffer ``flags`` marks current, on + every rank of the group, as the shard store of the one-shot kernel would place it; the + all-reduce of this output must then skip its own shard load and store. + """ + m, n, k = _validate_runtime_tensors(a, b, None, out) + validate_tactic(tactic, m, n, k) + if not publish_supported(m, n, tactic, dtype=a.dtype): + raise ValueError(f"Lamport publish does not serve M={m}, N={n} with {tactic}") + if ( + lamport.dtype != a.dtype + or lamport.dim() != 1 + or flags.dtype != _torch.int32 + or flags.dim() != 1 + ): + raise ValueError("lamport must be a 1-D view of the output dtype, flags a 1-D int32 view") + cute_tensors = _to_cute_swap(a, b, out, None) + compiled = _get_compiled_splitk_kernel( + dtype=a.dtype, + tactic=tactic, + use_pdl=pdl, + has_bias=False, + leading_dims=cute_tensors[4], + reserve_full_smem=reserve_full_smem, + publish=publish, + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + compiled( + *cute_tensors[:3], + _from_dlpack_dynamic(lamport, 0), + _from_dlpack_dynamic(flags, 0), + stream, + ) + + +@functools.cache +def _get_compiled_hidden_half_kernel( + tactic: SplitKTactic, + use_pdl: bool, + leading_dims: tuple[int, int, int], + norm_eps: float | None, +): + """The split-K GEMM with an FP32 C; with ``norm_eps``, the ``HnormB`` variant that takes the + norm weight as a fourth operand.""" + hnorm_b = None if norm_eps is None else HnormB(HNORM_B_HIDDEN, norm_eps) + kernel = SplitKDenseGemmKernel(tactic=tactic, use_pdl=use_pdl, has_bias=False, hnorm_b=hnorm_b) + # The representative tensors of _make_compile_repr_tensors, with an FP32 C. + m, n, k, batch = 64, 8, _CTA_K, 1 + compile_tensors = tuple( + _from_dlpack_dynamic(_make_layout_tensor(shape, dtype, leading_dim), leading_dim) + for shape, dtype, leading_dim in ( + ((batch, n, k), _torch.bfloat16, leading_dims[0]), + ((batch, k, m), _torch.bfloat16, leading_dims[1]), + ((batch, n, m), _torch.float32, leading_dims[2]), + ) + ) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + if hnorm_b is None: + return cute_ext.compile(_bmm_no_bias, kernel, *compile_tensors, stream) + norm_weight = _from_dlpack_dynamic(_torch.empty((k,), dtype=_torch.bfloat16, device="cuda"), 0) + return cute_ext.compile(_bmm_hnorm_b, kernel, *compile_tensors, norm_weight, stream) + + +def hidden_half_norm_supported( + x: _torch.Tensor, weight: _torch.Tensor, norm_weight: _torch.Tensor, tactic: SplitKTactic +) -> bool: + """Whether ``run_splitk_hidden_half`` can normalize ``x`` itself (``HnormB``): one 16-byte + aligned BF16 row of ``HNORM_B_HIDDEN`` with unit stride, a contiguous 16-byte aligned BF16 + norm weight of that size, a weight of that K, and a tactic of ``hnorm_b_supported``.""" + return ( + x.dim() == 2 + and x.shape[0] == 1 + and x.dtype == _torch.bfloat16 + and x.shape[1] == HNORM_B_HIDDEN + and x.stride(1) == 1 + and x.data_ptr() % 16 == 0 + and norm_weight.dtype == _torch.bfloat16 + and tuple(norm_weight.shape) == (x.shape[1],) + and norm_weight.is_contiguous() + and norm_weight.data_ptr() % 16 == 0 + and weight.dim() == 2 + and weight.shape[1] == x.shape[1] + and hnorm_b_supported(tactic, x.shape[1]) + ) + + +def run_splitk_hidden_half( + x: _torch.Tensor, + weight: _torch.Tensor, + out: _torch.Tensor, + pdl: bool, + tactic: SplitKTactic, + norm_weight: _torch.Tensor | None = None, + norm_eps: float | None = None, +) -> _torch.Tensor: + """``x @ weight.T`` into an FP32 ``out``, or with ``norm_weight`` ``rmsnorm(x) @ weight.T`` + with the norm computed in the kernel (``HnormB``). + + Args: + x: Activation [rows, K] BF16 with unit inner stride; with ``norm_weight`` the one raw + row whose FlashInfer RMSNorm (``norm_weight``, ``norm_eps``) the kernel computes. + weight: [N, K] BF16, contiguous. + out: [rows, N] FP32, contiguous: the split-K FP32 sum, never rounded to BF16. + pdl: Whether to use Programmatic Dependent Launch. + tactic: The ``SplitKTactic``; with ``norm_weight`` one of ``hnorm_b_supported``. + norm_weight: [K] BF16 RMSNorm weight, or None for an ``x`` that is already normalized. + norm_eps: The RMSNorm epsilon (with ``norm_weight``). + + Returns: + ``out``. + + With ``norm_weight`` the result is bit-identical to running FlashInfer's RMSNorm first and + this GEMM on its BF16 output. + """ + if x.dim() != 2 or weight.dim() != 2 or out.dim() != 2: + raise ValueError("x, weight and out must be 2D") + rows, k = x.shape + n = weight.shape[0] + if ( + x.dtype != _torch.bfloat16 + or weight.dtype != _torch.bfloat16 + or out.dtype != _torch.float32 + or weight.shape[1] != k + or tuple(out.shape) != (rows, n) + or not weight.is_contiguous() + or not out.is_contiguous() + or x.stride(1) != 1 + or any(t.data_ptr() % 32 for t in (x, weight, out)) + ): + raise ValueError( + "run_splitk_hidden_half needs BF16 x [rows, K] with unit inner stride, a contiguous " + "BF16 weight [N, K] and a contiguous FP32 out [rows, N], all 32-byte aligned; got " + f"x {tuple(x.shape)} {x.dtype}, weight {tuple(weight.shape)} {weight.dtype}, " + f"out {tuple(out.shape)} {out.dtype}" + ) + validate_tactic(tactic, rows, n, k) + if norm_weight is not None and ( + norm_eps is None or not hidden_half_norm_supported(x, weight, norm_weight, tactic) + ): + raise ValueError( + f"the in-kernel norm does not serve x {tuple(x.shape)}, weight " + f"{tuple(weight.shape)} with {tactic}" + ) + cute_tensors = _to_cute_swap(x, weight.t(), out, None) + compiled = _get_compiled_hidden_half_kernel( + tactic, pdl, cute_tensors[4], None if norm_weight is None else float(norm_eps) + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(x.device).cuda_stream) + if norm_weight is None: + compiled(*cute_tensors[:3], stream) + else: + compiled(*cute_tensors[:3], _from_dlpack_dynamic(norm_weight.detach(), 0), stream) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk_fused.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk_fused.py new file mode 100644 index 000000000000..4f7a258f7ebf --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk_fused.py @@ -0,0 +1,2033 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-M split-K BF16 GEMM with fused epilogues for the MLA q_b projection (SM10x, CuTe DSL). + +``SplitKFusedEpilogueGemmKernel`` runs the mainloop and the cluster split-K reduction of +``_low_m_splitk.SplitKDenseGemmKernel`` unchanged. Only the owner CTA's epilogue differs: +it stages the reduced FP32 tile in shared memory and then stores four consecutive output +columns per thread, so per-row and per-column-pair operations do not depend on the TMEM load +layout. Options (``SplitKEpilogue``): + +``row_rms`` + B (the public activation ``a``) is the raw input x of an RMSNorm whose weight is folded + into the GEMM weight (``W * gamma``). The owner's epilogue warps compute every row's mean + square from global memory while the mainloop runs, and the row's outputs are scaled by + ``rsqrt(mean + eps)``: ``(x * rstd) @ (W * gamma)^T``, which is ``RMSNorm(x) @ W^T`` up to + rounding. +``rms_out`` + Also write ``RMSNorm(x) = x * rstd * gamma`` (BF16 [M, K]) for other readers of the + normalized activation. The owner CTAs of a row tile split its K columns. +``fp8_rope`` + Write FP8 E4M3 (times a device scale) instead of BF16, after GPT-J RoPE on the last + ``rope_dim`` columns of every ``head_dim``-wide head: the MLA generation Q layout + [tokens, heads, kv_lora_rank + qk_rope_head_dim] that ``mla_rope_generation`` rotates + and quantizes, with its position rule (``kv_lens[seq] - seq_len + token_in_seq``), its + cos/sin table layout and its saturating E4M3 conversion. The BF16 rounding the unfused + path applies before the rotation and before the quantization is skipped. + +``SplitKIndexerQTailGemmKernel`` (``run_splitk_idx_qtail``) is the ``row_rms`` kernel of the DSA +indexer wq_b whose owner epilogue finishes the indexer Q (``IndexerQTail``) instead of storing it: +it rounds every output to BF16 as the BF16 store does, then applies the per-(token, head) tail of +the fused DSA indexer prologue to those values with the same operations in the same order (RoPE, +BF16 rounding, the per-head E4M3 scale, the E4M3 Q, the scale, and the attention weights times +the scale). Head-wide tiles finish a head per warp; half-head tiles run +in clusters of the two halves, whose owners exchange their row amax through DSMEM, so the GEMM +keeps its 64-row tactic and its K split. +""" + +from __future__ import annotations + +import dataclasses +import functools +import math + +import cuda.bindings.driver as _cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import torch as _torch +from cutlass import Float32, Int32, Int64 +from cutlass._mlir.dialects import llvm +from cutlass.cute import experimental as cute_ext +from cutlass.cutlass_dsl import T, dsl_user_op + +from ._low_m_splitk import ( + _AB_BUFFER_ALIGN_BYTES, + _CTA_K, + _LAUNCH_SMEM_HEADROOM_BYTES, + _MAILBOX_ALIGN_BYTES, + _MBARRIER_BYTES, + _SMEM_CAPACITY_BYTES, + _TMEM_POINTER_BYTES, + OWNER_RANK, + SplitKDenseGemmKernel, + SplitKTactic, + _align_up, + _detect_leading_dim, + _from_dlpack_dynamic, + _smem_bytes, + _store_shared_remote_v4, + default_tactic, + device_limits, + splitk_default_tactic, + validate_tactic, +) + +__all__ = [ + "IndexerQTail", + "Q_B_PLANS", + "SplitKEpilogue", + "fused_epilogue_tactic", + "idx_qtail_plan", + "idx_qtail_tactic", + "q_b_plan", + "run_splitk_fused", + "run_splitk_idx_qtail", +] + +#: Alignment of the epilogue's shared-memory staging buffers. +_EPI_ALIGN_BYTES = 16 + +#: Named barrier of the four epilogue warps (barrier 0 is the CTA-wide one). +_EPI_BARRIER_ID = 1 + +#: BF16 elements per 16-byte vector. +_BF16_PER_VEC = 8 + +#: Lanes of a warp. +_WARP_LANES = 32 + + +@dataclasses.dataclass(frozen=True) +class SplitKEpilogue: + """Fused epilogue options; every field is a compile-time constant of the kernel.""" + + row_rms: bool = False + rms_k: int = 0 + rms_eps: float = 0.0 + rms_out: bool = False + fp8_rope: bool = False + head_dim: int = 0 + rope_dim: int = 0 + seq_len: int = 1 + + def validate(self) -> None: + if not (self.row_rms or self.fp8_rope): + raise ValueError("a fused epilogue needs row_rms or fp8_rope") + if self.row_rms and (self.rms_k <= 0 or self.rms_k % _BF16_PER_VEC or self.rms_eps <= 0): + raise ValueError(f"row_rms needs a positive rms_k divisible by 8 and eps > 0: {self}") + if self.rms_out and not self.row_rms: + raise ValueError("rms_out needs row_rms") + if self.fp8_rope and not ( + 0 < self.rope_dim < self.head_dim + and self.rope_dim % 4 == 0 + and (self.head_dim - self.rope_dim) % 4 == 0 + and self.seq_len >= 1 + ): + raise ValueError(f"invalid fp8_rope geometry: {self}") + + +def _fused_smem_bytes(tactic: SplitKTactic, ab_stages: int) -> int: + """Shared memory of the fused kernel: the base layout, then the FP32 tile and row scales.""" + cursor = _align_up(_smem_bytes(tactic, ab_stages), _EPI_ALIGN_BYTES) + cursor += tactic.mma_m * tactic.mma_n * 4 + return _align_up(cursor, _EPI_ALIGN_BYTES) + tactic.mma_n * 4 + + +def fused_epilogue_tactic(tactic: SplitKTactic) -> SplitKTactic: + """``tactic`` with at most as many A/B stages as fit beside the epilogue staging buffers.""" + stages = tactic.ab_stages + while stages > 1 and _fused_smem_bytes(tactic, stages) > _SMEM_CAPACITY_BYTES: + stages -= 1 + return dataclasses.replace(tactic, ab_stages=stages) + + +@dataclasses.dataclass(frozen=True) +class IndexerQTail: + """The DSA indexer Q tail of ``run_splitk_idx_qtail``; every field is a compile-time constant. + + Per (token, head) of the BF16-rounded wq_b output, as ``_dsa_indexer_prologue_fused_kernel``: + RoPE on the first ``rope_dim`` columns (GPT-J pairs with ``interleave``, NeoX halves otherwise) + with the cos/sin row of the token's position, rounded to BF16; the head's E4M3 scale + max(amax, 1e-12) / 448, rounded up to a power of two with ``ue8m0``; the E4M3 Q (times the + MUFU reciprocal of the scale); the scale; and the attention weight times the scale times + ``weight_scale``. + """ + + n_heads: int + head_dim: int = 128 + rope_dim: int = 64 + interleave: bool = True + ue8m0: bool = True + weight_scale: float = 1.0 + + def validate(self) -> None: + if not ( + self.n_heads > 0 + and self.head_dim % 8 == 0 + and 0 < self.rope_dim <= self.head_dim // 2 + and self.rope_dim % 8 == 0 + and self.weight_scale > 0 + ): + raise ValueError(f"invalid indexer Q tail geometry: {self}") + + +def _qtail_pair(tactic: SplitKTactic, qtail: IndexerQTail) -> bool: + """Whether a tile holds half a head (clusters of the two halves) rather than a whole head.""" + if tactic.mma_m == qtail.head_dim: + return False + if 2 * tactic.mma_m == qtail.head_dim: + return True + raise ValueError(f"tile of {tactic.mma_m} columns is neither a head nor half a head: {qtail}") + + +def _qtail_smem_bytes(tactic: SplitKTactic, ab_stages: int, qtail: IndexerQTail) -> int: + """Shared memory of the indexer Q tail kernel: the fused layout, then the cos/sin rows and the + attention weights of the tile's tokens, then (half-head tiles) both halves' row amax and their + mbarrier.""" + cursor = _align_up(_fused_smem_bytes(tactic, ab_stages), _EPI_ALIGN_BYTES) + cursor += tactic.mma_n * qtail.rope_dim * 4 + cursor = _align_up(cursor, _EPI_ALIGN_BYTES) + tactic.mma_n * 4 + if _qtail_pair(tactic, qtail): + cursor = _align_up(cursor, _EPI_ALIGN_BYTES) + tactic.mma_n * 4 + cursor = _align_up(cursor, _EPI_ALIGN_BYTES) + tactic.mma_n * 4 + cursor = _align_up(cursor, _MBARRIER_BYTES) + _MBARRIER_BYTES + return cursor + + +def idx_qtail_tactic(tactic: SplitKTactic, qtail: IndexerQTail) -> SplitKTactic: + """``tactic`` with at most as many A/B stages as fit beside the indexer Q tail's buffers.""" + if _qtail_pair(tactic, qtail) and 2 * tactic.split_k > 8: + raise ValueError(f"half-head tiles need a cluster of at most 8 CTAs: {tactic}") + stages = tactic.ab_stages + while stages > 1 and _qtail_smem_bytes(tactic, stages, qtail) > _SMEM_CAPACITY_BYTES: + stages -= 1 + return dataclasses.replace(tactic, ab_stages=stages) + + +#: A/B pipeline stages of the absorbed q_b queued beside the DSA indexer top-k. +_BESIDE_TOPK_AB_STAGES = 3 + +#: The launch plans of the fused q_b (``q_b_plan``). +Q_B_PLANS = ("regular", "beside_topk", "early_gather") + + +def q_b_plan(plan: str, m: int, n: int, k: int, device_index: int) -> SplitKTactic: + """The split-K tactic of the fused q_b ``[m, k] @ [k, n]`` under ``plan``. + + ``regular``: ``splitk_default_tactic``. ``beside_topk``: 128-row tiles, no split-K and three + A/B stages, so that a q_b queued beside the DSA indexer top-k takes few SMs and little memory + bandwidth from it; the regular tactic where that one does not serve the shape. + ``early_gather``: the source default (split-K 4) where the regular tactic is split-K 8, which + leaves the clusters of a flash decode that gathers before its grid dependency wait room to + land while q_b runs. ``run_splitk_fused`` then lowers the A/B stages to fit its buffers. + """ + if plan not in Q_B_PLANS: + raise ValueError(f"unknown q_b plan {plan!r}; expected one of {Q_B_PLANS}") + sm_count, smem_capacity = device_limits(device_index) + tactic = splitk_default_tactic(m, n, k, sm_count=sm_count, smem_capacity=smem_capacity) + if plan == "beside_topk": + candidate = dataclasses.replace( + tactic, mma_m=128, split_k=1, ab_stages=_BESIDE_TOPK_AB_STAGES + ) + try: + validate_tactic(candidate, m, n, k) + except ValueError: + return tactic + return candidate + if plan == "early_gather" and tactic.split_k == 8: + return default_tactic(m, n, k) + return tactic + + +def idx_qtail_plan(qtail: IndexerQTail, m: int, n: int, k: int, device_index: int) -> SplitKTactic: + """The split-K tactic of the indexer wq_b ``[m, k] @ [k, n]`` with the Q tail: the + ``splitk_default_tactic`` of the shape, its 64-row tiles as half-head tiles, its A/B stages + lowered to fit the Q tail's buffers. Raises ``ValueError`` where it does not serve the shape. + """ + sm_count, smem_capacity = device_limits(device_index) + tactic = splitk_default_tactic(m, n, k, sm_count=sm_count, smem_capacity=smem_capacity) + tactic = idx_qtail_tactic(tactic, qtail) + validate_tactic(tactic, m, n, k) + return tactic + + +# ----------------------------------------------------------------------------- PTX helpers + + +@dsl_user_op +def _ld_global_v4(addr: Int64, *, loc=None, ip=None): + """16-byte global load as four 32-bit words.""" + res = llvm.inline_asm( + llvm.StructType.get_literal([T.i32(), T.i32(), T.i32(), T.i32()]), + [addr.ir_value(loc=loc, ip=ip)], + "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + return [Int32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(4)] + + +@dsl_user_op +def _bf16x2_to_f32x2(word: Int32, *, loc=None, ip=None): + """The two BF16 halves of a 32-bit word (low half first) as FP32.""" + res = llvm.inline_asm( + llvm.StructType.get_literal([T.f32(), T.f32()]), + [word.ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 bflo, bfhi;\n\tshl.b32 bflo, $2, 16;\n\tand.b32 bfhi, $2, 0xFFFF0000;\n\t" + "mov.b32 $0, bflo;\n\tmov.b32 $1, bfhi;\n\t}", + "=f,=f,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + return [Float32(llvm.extractvalue(T.f32(), res, [i], loc=loc, ip=ip)) for i in range(2)] + + +@dsl_user_op +def _cvt_bf16x2(lo: Float32, hi: Float32, *, loc=None, ip=None) -> Int32: + """Round two FP32 values to BF16 (nearest even), ``lo`` in the low half.""" + return Int32( + llvm.inline_asm( + T.i32(), + [lo.ir_value(loc=loc, ip=ip), hi.ir_value(loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $2, $1;", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _cvt_e4m3x4(v0: Float32, v1: Float32, v2: Float32, v3: Float32, *, loc=None, ip=None) -> Int32: + """Four FP32 values to saturated E4M3 (nearest even), ``v0`` in the lowest byte.""" + return Int32( + llvm.inline_asm( + T.i32(), + [v.ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "{\n\t.reg .b16 fp8lo, fp8hi;\n\tcvt.rn.satfinite.e4m3x2.f32 fp8lo, $2, $1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 fp8hi, $4, $3;\n\tmov.b32 $0, {fp8lo, fp8hi};\n\t}", + "=r,f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _st_global_b32(addr: Int64, v: Int32, *, loc=None, ip=None) -> None: + llvm.inline_asm( + None, + [addr.ir_value(loc=loc, ip=ip), v.ir_value(loc=loc, ip=ip)], + "st.global.b32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _st_global_v2(addr: Int64, v0: Int32, v1: Int32, *, loc=None, ip=None) -> None: + llvm.inline_asm( + None, + [addr.ir_value(loc=loc, ip=ip), v0.ir_value(loc=loc, ip=ip), v1.ir_value(loc=loc, ip=ip)], + "st.global.v2.b32 [$0], {$1, $2};", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _st_global_v4( + addr: Int64, v0: Int32, v1: Int32, v2: Int32, v3: Int32, *, loc=None, ip=None +) -> None: + llvm.inline_asm( + None, + [addr.ir_value(loc=loc, ip=ip)] + [v.ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "st.global.v4.b32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _fma_rn(a: Float32, b: Float32, c: Float32, *, loc=None, ip=None) -> Float32: + """``a * b + c`` rounded once (nearest even).""" + return Float32( + llvm.inline_asm( + T.f32(), + [v.ir_value(loc=loc, ip=ip) for v in (a, b, c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _mul_rn(a: Float32, b: Float32, *, loc=None, ip=None) -> Float32: + """``a * b`` rounded once (nearest even), never contracted into an FMA.""" + return Float32( + llvm.inline_asm( + T.f32(), + [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _neg_f32(a: Float32, *, loc=None, ip=None) -> Float32: + """``-a`` (sign flip, exact).""" + return Float32( + llvm.inline_asm( + T.f32(), + [a.ir_value(loc=loc, ip=ip)], + "neg.f32 $0, $1;", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _fp8_ratio(amax: Float32, *, loc=None, ip=None) -> Float32: + """max(amax, 1e-12) * fp32(1 / 448): the per-row E4M3 ratio of the indexer prologue + (``_fp8_row_scale``), with its FP32 constants.""" + return Float32( + llvm.inline_asm( + T.f32(), + [amax.ir_value(loc=loc, ip=ip)], + "{\n\t.reg .f32 fpfloor;\n\tmax.f32 fpfloor, $1, 0f2B8CBCCC;\n\t" + "mul.rn.f32 $0, fpfloor, 0f3B124925;\n\t}", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _ue8m0_up(ratio: Float32, *, loc=None, ip=None) -> Float32: + """The positive ``ratio`` with its mantissa cleared and, if the mantissa was nonzero, its + exponent raised by one: the UE8M0 scale of the indexer prologue.""" + return Float32( + llvm.inline_asm( + T.f32(), + [ratio.ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 uebits, uemant, ueexp;\n\t.reg .pred uenz;\n\tmov.b32 uebits, $1;\n\t" + "and.b32 uemant, uebits, 0x007FFFFF;\n\tand.b32 ueexp, uebits, 0x7F800000;\n\t" + "setp.ne.u32 uenz, uemant, 0;\n\t@uenz add.u32 ueexp, ueexp, 0x00800000;\n\t" + "mov.b32 $0, ueexp;\n\t}", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +def _bf16_round(lo: Float32, hi: Float32): + """Two FP32 values rounded to BF16 (nearest even), back as FP32.""" + return _bf16x2_to_f32x2(_cvt_bf16x2(lo, hi)) + + +# --------------------------------------------------------------------------------- kernel + + +class SplitKFusedEpilogueGemmKernel(SplitKDenseGemmKernel): + """``SplitKDenseGemmKernel`` (no bias) with a fused owner epilogue.""" + + def __init__( + self, + *, + tactic: SplitKTactic, + use_pdl: bool, + epilogue: SplitKEpilogue, + reserve_full_smem: bool = True, + ) -> None: + super().__init__( + tactic=tactic, use_pdl=use_pdl, has_bias=False, reserve_full_smem=reserve_full_smem + ) + epilogue.validate() + self.epi = epilogue + self.launch_smem_bytes = min( + _align_up(_fused_smem_bytes(tactic, tactic.ab_stages), 1024) + + _LAUNCH_SMEM_HEADROOM_BYTES, + _SMEM_CAPACITY_BYTES, + ) + if (self.cta_m * self.cta_n // 4) % self.epilog_threads: + raise ValueError( + f"tile ({self.cta_m}, {self.cta_n}) does not split into 4-column groups" + ) + self.rows_per_tile = self.cta_n + self.threads_per_row = self.epilog_threads // self.cta_n + if epilogue.row_rms and ( + self.threads_per_row & (self.threads_per_row - 1) + or self.threads_per_row > 32 + or (epilogue.rms_k // _BF16_PER_VEC) % self.threads_per_row + ): + raise ValueError( + f"row_rms cannot split K={epilogue.rms_k} over a ({self.cta_m}, {self.cta_n}) tile" + ) + + @cute.experimental.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + gamma: cute.Tensor, + norm_out: cute.Tensor, + kv_lens: cute.Tensor, + cos_sin: cute.Tensor, + q_scale: cute.Tensor, + stream: _cuda.CUstream, + ): + self.kernel(a, b, c, gamma, norm_out, kv_lens, cos_sin, q_scale).launch( + grid=( + cute.ceil_div(c.layout.shape[0], self.cta_m), + cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k, + c.layout.shape[2], + ), + block=(self.threads_per_cta, 1, 1), + cluster=self.cluster_shape, + smem=cute.Int64( + utils.get_smem_capacity_in_bytes("sm_100") + if self.reserve_full_smem + else self.launch_smem_bytes + ), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.experimental.kernel + def kernel( + self, + mA: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L), K-major: the weight + mB: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L), K-major: the activation (rows = tokens) + mC: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L), M-major: BF16, or uint8 for fp8_rope + mGamma: cute.Tensor, # [K] BF16 norm weight (rms_out) + mNormOut: cute.Tensor, # [M, K] BF16 normalized activation (rms_out) + mKvLens: cute.Tensor, # [sequences] int32 KV lengths (fp8_rope) + mCosSin: cute.Tensor, # flat FP32 (cos, sin) table (fp8_rope) + mQScale: cute.Tensor, # [1] FP32 quantization scale (fp8_rope) + ): + """Allocate storage and dispatch the specialized warps (as the base kernel).""" + stages = self.num_ab_stage + + ab_dtype = mA.element_type + tiled_mma = sm100_utils.make_trivial_tiled_mma( + ab_dtype, + ab_dtype, + utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + + mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k) + block_idx = cute.arch.block_idx() + bidx = block_idx[0] + split_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + n_idx = block_idx[1] // self.split_k + l_idx = block_idx[2] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + sA = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + sB = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + + acc_layout = cute_ext.make_tmem_layout_acc(tiled_mma, self.mma_tiler_mn, acc_stage=1) + c_tiler_mn = (self.cta_m, self.cta_n) + + bar_full = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_empty = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tma_epilog = cute_ext.allocate( + cutlass.Int64, cute.AddressSpace.smem, cute.make_layout(1), alignment=_MBARRIER_BYTES + ).iterator + bar_mma_epilog = cute_ext.allocate( + cutlass.Int64, cute.AddressSpace.smem, cute.make_layout(1), alignment=_MBARRIER_BYTES + ).iterator + bar_tmem_alloc = cute_ext.allocate( + cutlass.Int64, cute.AddressSpace.smem, cute.make_layout(1), alignment=_MBARRIER_BYTES + ).iterator + tmem_base_ptr = cute_ext.allocate( + cutlass.Int32, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_TMEM_POINTER_BYTES, + ).iterator + + if cutlass.const_expr(self.split_k > 1): + mailbox = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.mailbox_elements), + alignment=_MAILBOX_ALIGN_BYTES, + ) + bar_reduce = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + else: + mailbox = sA + bar_reduce = bar_mma_epilog + + # Epilogue staging: the reduced FP32 tile (M-major, as the output tile) and the + # per-row scales. + sAcc = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout((self.cta_m, self.cta_n)), + alignment=_EPI_ALIGN_BYTES, + ) + sRstd = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.cta_n), + alignment=_EPI_ALIGN_BYTES, + ) + + if warp_idx == 0: + with cute.arch.elect_one(): + for i in range(stages): + cute.arch.mbarrier_init(bar_full + i, 2) + cute.arch.mbarrier_init(bar_empty + i, 1) + cute.arch.mbarrier_init(bar_tma_epilog, 32) + cute.arch.mbarrier_init(bar_mma_epilog, 1) + cute.arch.mbarrier_init(bar_tmem_alloc, 160) + + if cutlass.const_expr(self.split_k > 1): + cute.arch.mbarrier_init(bar_reduce, 1) + + cute.arch.mbarrier_init_fence() + if cutlass.const_expr(self.split_k > 1): + cute.arch.cluster_arrive_relaxed() + else: + cute.arch.barrier() + + k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k + k_tile_start = split_rank * k_tile_count + + if cutlass.const_expr(self.split_k > 1): + cute.arch.cluster_wait() + + if warp_idx == 0: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), + sA, + cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), + k_tile_start, + k_tile_count, + True, + ) + elif warp_idx == 1: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), + sB, + cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), + k_tile_start, + k_tile_count, + False, + ) + elif warp_idx == 2: + self.mma_warp( + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma, + sA, + sB, + tmem_base_ptr, + acc_layout, + self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), + k_tile_count, + ) + elif warp_idx >= 4: + self.epilog_fused( + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout, + cute.local_tile(mC, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.arch.thread_idx()[0] - 128, + utils.LayoutEnum.from_tensor(mC), + mailbox, + bar_reduce, + split_rank, + sAcc, + sRstd, + mB, + mC, + mGamma, + mNormOut, + mKvLens, + mCosSin, + mQScale, + bidx, + n_idx, + ) + + @cute.experimental.jit + def epilog_fused( + self, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + gD_tile: cute.Tensor, + epi_tid: cutlass.Int32, + d_layout: cutlass.Constexpr, + mailbox, + bar_reduce, + split_rank: cutlass.Int32, + sAcc: cute.Tensor, + sRstd: cute.Tensor, + mB: cute.Tensor, + mC: cute.Tensor, + mGamma: cute.Tensor, + mNormOut: cute.Tensor, + mKvLens: cute.Tensor, + mCosSin: cute.Tensor, + mQScale: cute.Tensor, + bidx: cutlass.Int32, + n_idx: cutlass.Int32, + ): + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) + + acc_view = cute.make_tensor( + cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr), + acc_layout, + )[((None, None), 0, 0, 0)] + + epi_tile = (self.cta_m, self.cta_n) + # The TMEM load op of the BF16 kernel, whatever the output type: the tile is staged + # in shared memory, so only the accumulator partition matters. + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( + sm100_utils.get_tmem_load_op( + (self.cta_m, self.cta_n, self.cta_k), + d_layout, + cutlass.BFloat16, + self.acc_dtype, + epi_tile, + False, + ), + acc_view, + ) + gD_epi = cute.flat_divide(gD_tile, epi_tile) + rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid) + rAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + thr_t2r = tiled_copy_t2r.get_slice(epi_tid) + + # Row statistics overlap the mainloop: only the owner needs them. + if cutlass.const_expr(self.epi.row_rms): + if split_rank == OWNER_RANK: + self._row_rstd(mB, sRstd, n_idx, epi_tid) + if cutlass.const_expr(self.epi.rms_out): + self._rms_out(mB, mC, mGamma, mNormOut, sRstd, bidx, n_idx, epi_tid) + + cute.arch.mbarrier_wait(bar_mma_epilog, 0) + cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) + cute.arch.fence_view_async_tmem_load() + cute.arch.mbarrier_arrive(bar_tmem_alloc) + + if cutlass.const_expr(self.split_k > 1): + values_per_thread = cutlass.const_expr(cute.size(rmem_layout)) + values_per_peer = cutlass.const_expr(self.epilog_threads * values_per_thread) + if split_rank != OWNER_RANK: + for value_idx in cutlass.range_constexpr(0, values_per_thread, 4): + _store_shared_remote_v4( + rAcc[value_idx], + rAcc[value_idx + 1], + rAcc[value_idx + 2], + rAcc[value_idx + 3], + mailbox.iterator + + (split_rank - Int32(1)) * values_per_peer + + epi_tid * values_per_thread + + value_idx, + bar_reduce, + Int32(OWNER_RANK), + ) + else: + if epi_tid == 0: + cute.arch.mbarrier_arrive_and_expect_tx( + bar_reduce, self.expected_transaction_bytes + ) + cute.arch.mbarrier_wait(bar_reduce, 0) + for peer in cutlass.range_constexpr(self.split_k - 1): + for value_idx in cutlass.range_constexpr(values_per_thread): + rAcc[value_idx] = ( + rAcc[value_idx] + + mailbox[ + peer * values_per_peer + epi_tid * values_per_thread + value_idx + ] + ) + + if split_rank == OWNER_RANK: + # Stage the tile by output coordinate, then store four columns per thread. + cute_ext.partition_and_copy(thr_t2r, rAcc, sAcc) + cute.arch.barrier(barrier_id=_EPI_BARRIER_ID, number_of_threads=self.epilog_threads) + self._store_tile(sAcc, sRstd, mC, mKvLens, mCosSin, mQScale, bidx, n_idx, epi_tid) + + @cute.experimental.jit + def _row_rstd( + self, + mB: cute.Tensor, + sRstd: cute.Tensor, + n_idx: cutlass.Int32, + epi_tid: cutlass.Int32, + ): + """sRstd[row] = rsqrt(mean(x[row]^2) + eps) for the tile's rows; ends with a barrier.""" + tpr = cutlass.const_expr(self.threads_per_row) + vecs_per_thread = cutlass.const_expr(self.epi.rms_k // _BF16_PER_VEC // tpr) + # The activation is the previous grid's output. + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + row = epi_tid // tpr + lane = epi_tid - row * tpr + token = n_idx * self.cta_n + row + ssq = Float32(0.0) + if token < cute.size(mB, mode=[0]): + base = (mB.iterator + cute.crd2idx((token, 0, 0), mB.layout)).toint() + Int64(lane * 16) + for i in cutlass.range_constexpr(vecs_per_thread): + words = _ld_global_v4(base + Int64(i * tpr * 16)) + for j in cutlass.range_constexpr(4): + pair = _bf16x2_to_f32x2(words[j]) + ssq = ssq + pair[0] * pair[0] + pair[1] * pair[1] + for i in cutlass.range_constexpr(int(math.log2(tpr))): + ssq = ssq + cute.arch.shuffle_sync_bfly(ssq, offset=1 << i) + if lane == 0: + sRstd[row] = cute.math.rsqrt( + ssq / Float32(self.epi.rms_k) + Float32(self.epi.rms_eps), fastmath=True + ) + cute.arch.barrier(barrier_id=_EPI_BARRIER_ID, number_of_threads=self.epilog_threads) + + @cute.experimental.jit + def _rms_out( + self, + mB: cute.Tensor, + mC: cute.Tensor, + mGamma: cute.Tensor, + mNormOut: cute.Tensor, + sRstd: cute.Tensor, + bidx: cutlass.Int32, + n_idx: cutlass.Int32, + epi_tid: cutlass.Int32, + ): + """Write x * rstd * gamma for this tile's rows and its share of the K columns.""" + vecs = cutlass.const_expr(self.epi.rms_k // _BF16_PER_VEC) + n_tiles = cute.ceil_div(cute.size(mC, mode=[0]), self.cta_m) + chunk = (vecs + n_tiles - 1) // n_tiles + total = chunk * self.cta_n + iters = (total + self.epilog_threads - 1) // self.epilog_threads + tokens = cute.size(mB, mode=[0]) + for i in cutlass.range(iters, unroll=1): + idx = epi_tid + i * self.epilog_threads + row = idx // chunk + vec = bidx * chunk + (idx - row * chunk) + token = n_idx * self.cta_n + row + if idx < total: + if vec < vecs: + if token < tokens: + col = vec * _BF16_PER_VEC + x_words = _ld_global_v4( + (mB.iterator + cute.crd2idx((token, col, 0), mB.layout)).toint() + ) + g_words = _ld_global_v4((mGamma.iterator + col).toint()) + rstd = sRstd[row] + out_words = [] + for j in cutlass.range_constexpr(4): + x_pair = _bf16x2_to_f32x2(x_words[j]) + g_pair = _bf16x2_to_f32x2(g_words[j]) + out_words.append( + _cvt_bf16x2( + x_pair[0] * rstd * g_pair[0], x_pair[1] * rstd * g_pair[1] + ) + ) + out_ptr = mNormOut.iterator + cute.crd2idx((token, col), mNormOut.layout) + _st_global_v4( + out_ptr.toint(), + out_words[0], + out_words[1], + out_words[2], + out_words[3], + ) + + @cute.experimental.jit + def _store_tile( + self, + sAcc: cute.Tensor, + sRstd: cute.Tensor, + mC: cute.Tensor, + mKvLens: cute.Tensor, + mCosSin: cute.Tensor, + mQScale: cute.Tensor, + bidx: cutlass.Int32, + n_idx: cutlass.Int32, + epi_tid: cutlass.Int32, + ): + """Store the staged tile: four consecutive output columns of one row per thread.""" + groups_m = cutlass.const_expr(self.cta_m // 4) + groups_per_thread = cutlass.const_expr(self.cta_m * self.cta_n // 4 // self.epilog_threads) + n_feat = cute.size(mC, mode=[0]) + n_tok = cute.size(mC, mode=[1]) + q_scale = Float32(1.0) + if cutlass.const_expr(self.epi.fp8_rope): + q_scale = mQScale[0] + for g in cutlass.range_constexpr(groups_per_thread): + gi = epi_tid + g * self.epilog_threads + row = gi // groups_m + m4 = (gi - row * groups_m) * 4 + v0 = sAcc[m4, row] + v1 = sAcc[m4 + 1, row] + v2 = sAcc[m4 + 2, row] + v3 = sAcc[m4 + 3, row] + if cutlass.const_expr(self.epi.row_rms): + rstd = sRstd[row] + v0 = v0 * rstd + v1 = v1 * rstd + v2 = v2 * rstd + v3 = v3 * rstd + token = n_idx * self.cta_n + row + feature = bidx * self.cta_m + m4 + if token < n_tok: + if feature < n_feat: + addr = (mC.iterator + cute.crd2idx((feature, token, 0), mC.layout)).toint() + if cutlass.const_expr(self.epi.fp8_rope): + nope = cutlass.const_expr(self.epi.head_dim - self.epi.rope_dim) + col = feature % self.epi.head_dim + if col >= nope: + seq = token // self.epi.seq_len + pos = mKvLens[seq] - self.epi.seq_len + (token - seq * self.epi.seq_len) + # The table holds rope_dim (cos, sin) pairs per position; pair p + # rotates columns 2p and 2p + 1 of the head's RoPE slice. + cs = (pos * self.epi.rope_dim + (col - nope) // 2) * 2 + c0 = mCosSin[cs] + s0 = mCosSin[cs + 1] + c1 = mCosSin[cs + 2] + s1 = mCosSin[cs + 3] + r0 = c0 * v0 - s0 * v1 + r1 = c0 * v1 + s0 * v0 + r2 = c1 * v2 - s1 * v3 + r3 = c1 * v3 + s1 * v2 + v0 = r0 + v1 = r1 + v2 = r2 + v3 = r3 + _st_global_b32( + addr, + _cvt_e4m3x4(v0 * q_scale, v1 * q_scale, v2 * q_scale, v3 * q_scale), + ) + else: + _st_global_v2(addr, _cvt_bf16x2(v0, v1), _cvt_bf16x2(v2, v3)) + + +class SplitKIndexerQTailGemmKernel(SplitKFusedEpilogueGemmKernel): + """The ``row_rms`` kernel of the DSA indexer wq_b with the indexer Q tail (``IndexerQTail``) + in the owner epilogue, which writes the E4M3 Q, its per-head scale and the scaled attention + weights instead of the BF16 output. + + A tile is one head (``mma_m == head_dim``: a warp finishes one token's head) or half a head + (``2 * mma_m == head_dim``: half a warp per token, and the two halves of a head run in one + cluster of shape (2, split_k, 1), whose owners exchange their row amax through DSMEM). The + mainloop, the split-K reduction and the row RMS are those of the fused kernel, so the GEMM of + a tactic is unchanged. + """ + + def __init__( + self, + *, + tactic: SplitKTactic, + use_pdl: bool, + epilogue: SplitKEpilogue, + qtail: IndexerQTail, + reserve_full_smem: bool = True, + ) -> None: + super().__init__( + tactic=tactic, use_pdl=use_pdl, epilogue=epilogue, reserve_full_smem=reserve_full_smem + ) + qtail.validate() + if not epilogue.row_rms or epilogue.rms_out or epilogue.fp8_rope: + raise ValueError("the indexer Q tail follows the row_rms epilogue alone") + self.qtail = qtail + self.pair = _qtail_pair(tactic, qtail) + self.cluster_shape = (2 if self.pair else 1, tactic.split_k, 1) + self.lanes_per_row = self.cta_m // 4 + if _WARP_LANES % self.lanes_per_row or (self.epilog_threads // self.cta_n) == 0: + raise ValueError(f"tile ({self.cta_m}, {self.cta_n}) does not map rows onto lanes") + if qtail.rope_dim % (self.epilog_threads // self.cta_n): + raise ValueError(f"rope_dim {qtail.rope_dim} does not split over a tile row's threads") + self.launch_smem_bytes = min( + _align_up(_qtail_smem_bytes(tactic, tactic.ab_stages, qtail), 1024) + + _LAUNCH_SMEM_HEADROOM_BYTES, + _SMEM_CAPACITY_BYTES, + ) + + @cute.experimental.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + positions: cute.Tensor, + cos_sin: cute.Tensor, + weights_in: cute.Tensor, + q_scale_out: cute.Tensor, + weights_out: cute.Tensor, + stream: _cuda.CUstream, + ): + self.kernel(a, b, c, positions, cos_sin, weights_in, q_scale_out, weights_out).launch( + grid=( + cute.ceil_div(c.layout.shape[0], self.cta_m), + cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k, + c.layout.shape[2], + ), + block=(self.threads_per_cta, 1, 1), + cluster=self.cluster_shape, + smem=cute.Int64( + utils.get_smem_capacity_in_bytes("sm_100") + if self.reserve_full_smem + else self.launch_smem_bytes + ), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.experimental.kernel + def kernel( + self, + mA: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L), K-major: the weight + mB: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L), K-major: the activation (rows = tokens) + mC: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L), M-major: uint8 (E4M3 bits) + mPos: cute.Tensor, # [tokens] int32 positions + mCosSin: cute.Tensor, # [positions, >= rope_dim] FP32 [cos | sin] rows + mWIn: cute.Tensor, # [tokens, heads] BF16 attention weights (weights_proj output) + mQScaleOut: cute.Tensor, # [tokens, heads] FP32 E4M3 scales + mWOut: cute.Tensor, # [tokens, heads] FP32 attention weights * scale * weight_scale + ): + """Allocate storage and dispatch the specialized warps (as the fused kernel).""" + stages = self.num_ab_stage + + ab_dtype = mA.element_type + tiled_mma = sm100_utils.make_trivial_tiled_mma( + ab_dtype, + ab_dtype, + utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + + mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k) + block_idx = cute.arch.block_idx() + bidx = block_idx[0] + # The cluster's ranks: x is the half of the head (half-head tiles), y the K slice. + split_rank = cute.arch.make_warp_uniform(block_idx[1] % self.split_k) + half = Int32(0) + if cutlass.const_expr(self.pair): + half = cute.arch.make_warp_uniform(bidx % 2) + n_idx = block_idx[1] // self.split_k + l_idx = block_idx[2] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + sA = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + sB = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + + acc_layout = cute_ext.make_tmem_layout_acc(tiled_mma, self.mma_tiler_mn, acc_stage=1) + c_tiler_mn = (self.cta_m, self.cta_n) + + bar_full = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_empty = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tma_epilog = cute_ext.allocate( + cutlass.Int64, cute.AddressSpace.smem, cute.make_layout(1), alignment=_MBARRIER_BYTES + ).iterator + bar_mma_epilog = cute_ext.allocate( + cutlass.Int64, cute.AddressSpace.smem, cute.make_layout(1), alignment=_MBARRIER_BYTES + ).iterator + bar_tmem_alloc = cute_ext.allocate( + cutlass.Int64, cute.AddressSpace.smem, cute.make_layout(1), alignment=_MBARRIER_BYTES + ).iterator + tmem_base_ptr = cute_ext.allocate( + cutlass.Int32, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_TMEM_POINTER_BYTES, + ).iterator + + if cutlass.const_expr(self.split_k > 1): + mailbox = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.mailbox_elements), + alignment=_MAILBOX_ALIGN_BYTES, + ) + bar_reduce = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + else: + mailbox = sA + bar_reduce = bar_mma_epilog + + sAcc = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout((self.cta_m, self.cta_n)), + alignment=_EPI_ALIGN_BYTES, + ) + sRstd = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.cta_n), + alignment=_EPI_ALIGN_BYTES, + ) + # The tile's tokens: their cos/sin rows and attention weights, loaded while the mainloop + # runs. + sCS = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.cta_n * self.qtail.rope_dim), + alignment=_EPI_ALIGN_BYTES, + ) + sW = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.cta_n), + alignment=_EPI_ALIGN_BYTES, + ) + if cutlass.const_expr(self.pair): + # Row amax of this half, and of the other half (written by its owner). + sAmaxMine = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.cta_n), + alignment=_EPI_ALIGN_BYTES, + ) + sAmaxPeer = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.cta_n), + alignment=_EPI_ALIGN_BYTES, + ) + bar_amax = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + else: + sAmaxMine = sRstd + sAmaxPeer = sRstd + bar_amax = bar_mma_epilog + + if warp_idx == 0: + with cute.arch.elect_one(): + for i in range(stages): + cute.arch.mbarrier_init(bar_full + i, 2) + cute.arch.mbarrier_init(bar_empty + i, 1) + cute.arch.mbarrier_init(bar_tma_epilog, 32) + cute.arch.mbarrier_init(bar_mma_epilog, 1) + cute.arch.mbarrier_init(bar_tmem_alloc, 160) + + if cutlass.const_expr(self.split_k > 1): + cute.arch.mbarrier_init(bar_reduce, 1) + if cutlass.const_expr(self.pair): + # The local arrival plus the other owner's transaction bytes. + cute.arch.mbarrier_init(bar_amax, 1) + + cute.arch.mbarrier_init_fence() + if cutlass.const_expr(self.split_k > 1 or self.pair): + cute.arch.cluster_arrive_relaxed() + else: + cute.arch.barrier() + + k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k + k_tile_start = split_rank * k_tile_count + + if cutlass.const_expr(self.split_k > 1 or self.pair): + cute.arch.cluster_wait() + + if warp_idx == 0: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), + sA, + cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), + k_tile_start, + k_tile_count, + True, + ) + elif warp_idx == 1: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), + sB, + cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), + k_tile_start, + k_tile_count, + False, + ) + elif warp_idx == 2: + self.mma_warp( + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma, + sA, + sB, + tmem_base_ptr, + acc_layout, + self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), + k_tile_count, + ) + elif warp_idx >= 4: + self.epilog_qtail( + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout, + cute.local_tile(mC, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.arch.thread_idx()[0] - 128, + utils.LayoutEnum.from_tensor(mC), + mailbox, + bar_reduce, + split_rank, + half, + sAcc, + sRstd, + sCS, + sW, + sAmaxMine, + sAmaxPeer, + bar_amax, + mB, + mC, + mPos, + mCosSin, + mWIn, + mQScaleOut, + mWOut, + bidx, + n_idx, + ) + + @cute.experimental.jit + def epilog_qtail( + self, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + gD_tile: cute.Tensor, + epi_tid: cutlass.Int32, + d_layout: cutlass.Constexpr, + mailbox, + bar_reduce, + split_rank: cutlass.Int32, + half: cutlass.Int32, + sAcc: cute.Tensor, + sRstd: cute.Tensor, + sCS: cute.Tensor, + sW: cute.Tensor, + sAmaxMine: cute.Tensor, + sAmaxPeer: cute.Tensor, + bar_amax, + mB: cute.Tensor, + mC: cute.Tensor, + mPos: cute.Tensor, + mCosSin: cute.Tensor, + mWIn: cute.Tensor, + mQScaleOut: cute.Tensor, + mWOut: cute.Tensor, + bidx: cutlass.Int32, + n_idx: cutlass.Int32, + ): + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) + + acc_view = cute.make_tensor( + cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr), + acc_layout, + )[((None, None), 0, 0, 0)] + + epi_tile = (self.cta_m, self.cta_n) + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( + sm100_utils.get_tmem_load_op( + (self.cta_m, self.cta_n, self.cta_k), + d_layout, + cutlass.BFloat16, + self.acc_dtype, + epi_tile, + False, + ), + acc_view, + ) + gD_epi = cute.flat_divide(gD_tile, epi_tile) + rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid) + rAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + thr_t2r = tiled_copy_t2r.get_slice(epi_tid) + + head = bidx + if cutlass.const_expr(self.pair): + head = bidx // 2 + # The owner's row statistics, cos/sin rows and attention weights overlap the mainloop. + if split_rank == OWNER_RANK: + self._row_rstd(mB, sRstd, n_idx, epi_tid) + self._qtail_prefetch(mB, mPos, mCosSin, mWIn, sCS, sW, half, head, n_idx, epi_tid) + + cute.arch.mbarrier_wait(bar_mma_epilog, 0) + cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) + cute.arch.fence_view_async_tmem_load() + cute.arch.mbarrier_arrive(bar_tmem_alloc) + + if cutlass.const_expr(self.split_k > 1): + # The owner of this tile is the cluster's CTA of the same half at K slice 0. + owner_rank = half + values_per_thread = cutlass.const_expr(cute.size(rmem_layout)) + values_per_peer = cutlass.const_expr(self.epilog_threads * values_per_thread) + if split_rank != OWNER_RANK: + for value_idx in cutlass.range_constexpr(0, values_per_thread, 4): + _store_shared_remote_v4( + rAcc[value_idx], + rAcc[value_idx + 1], + rAcc[value_idx + 2], + rAcc[value_idx + 3], + mailbox.iterator + + (split_rank - Int32(1)) * values_per_peer + + epi_tid * values_per_thread + + value_idx, + bar_reduce, + owner_rank, + ) + else: + if epi_tid == 0: + cute.arch.mbarrier_arrive_and_expect_tx( + bar_reduce, self.expected_transaction_bytes + ) + cute.arch.mbarrier_wait(bar_reduce, 0) + for peer in cutlass.range_constexpr(self.split_k - 1): + for value_idx in cutlass.range_constexpr(values_per_thread): + rAcc[value_idx] = ( + rAcc[value_idx] + + mailbox[ + peer * values_per_peer + epi_tid * values_per_thread + value_idx + ] + ) + + if split_rank == OWNER_RANK: + cute_ext.partition_and_copy(thr_t2r, rAcc, sAcc) + cute.arch.barrier(barrier_id=_EPI_BARRIER_ID, number_of_threads=self.epilog_threads) + self._store_qtail( + sAcc, + sRstd, + sCS, + sW, + sAmaxMine, + sAmaxPeer, + bar_amax, + mC, + mQScaleOut, + mWOut, + bidx, + half, + head, + n_idx, + epi_tid, + ) + + @cute.experimental.jit + def _qtail_prefetch( + self, + mB: cute.Tensor, + mPos: cute.Tensor, + mCosSin: cute.Tensor, + mWIn: cute.Tensor, + sCS: cute.Tensor, + sW: cute.Tensor, + half: cutlass.Int32, + head: cutlass.Int32, + n_idx: cutlass.Int32, + epi_tid: cutlass.Int32, + ): + """sCS[row] = the cos/sin row of the token's position and sW[row] = its attention weight, + for the tile's tokens, in the tile that holds the head's RoPE columns. Runs after + ``_row_rstd``, whose grid dependency wait covers the positions and the weights.""" + rope = cutlass.const_expr(self.qtail.rope_dim) + tpr = cutlass.const_expr(self.epilog_threads // self.cta_n) + per = cutlass.const_expr(rope // tpr) + row = epi_tid // tpr + lane = epi_tid - row * tpr + token = n_idx * self.cta_n + row + if token < cute.size(mB, mode=[0]): + if half == 0: + pos = mPos[token] + for i in cutlass.range_constexpr(per): + col = lane * per + i + sCS[row * rope + col] = mCosSin[pos, col] + if lane == 0: + sW[row] = mWIn[token, head].to(Float32) + + @cute.experimental.jit + def _store_qtail( + self, + sAcc: cute.Tensor, + sRstd: cute.Tensor, + sCS: cute.Tensor, + sW: cute.Tensor, + sAmaxMine: cute.Tensor, + sAmaxPeer: cute.Tensor, + bar_amax, + mC: cute.Tensor, + mQScaleOut: cute.Tensor, + mWOut: cute.Tensor, + bidx: cutlass.Int32, + half: cutlass.Int32, + head: cutlass.Int32, + n_idx: cutlass.Int32, + epi_tid: cutlass.Int32, + ): + """Finish the staged tile: four consecutive columns of one token row per thread, the + lanes of a row reducing the head's amax with butterfly shuffles.""" + qt = self.qtail + rope = cutlass.const_expr(qt.rope_dim) + rhalf = cutlass.const_expr(qt.rope_dim // 2) + lanes = cutlass.const_expr(self.lanes_per_row) + groups_per_thread = cutlass.const_expr(self.cta_m * self.cta_n // 4 // self.epilog_threads) + n_tok = cute.size(mC, mode=[1]) + col0 = half * self.cta_m + rows = [] + m4s = [] + vals = [] + amaxes = [] + for g in cutlass.range_constexpr(groups_per_thread): + gi = epi_tid + g * self.epilog_threads + row = gi // lanes + m4 = (gi - row * lanes) * 4 + rstd = sRstd[row] + # The BF16 wq_b output, as the BF16 store rounds it. + x01 = _bf16_round(sAcc[m4, row] * rstd, sAcc[m4 + 1, row] * rstd) + x23 = _bf16_round(sAcc[m4 + 2, row] * rstd, sAcc[m4 + 3, row] * rstd) + q0 = x01[0] + q1 = x01[1] + q2 = x23[0] + q3 = x23[1] + col = col0 + m4 + if col < rope: + # x * cos + (+-partner) * sin as fma(x, cos, rn(+-partner * sin)), then BF16. + if cutlass.const_expr(qt.interleave): + cs = row * rope + col // 2 + c0 = sCS[cs] + s0 = sCS[cs + rhalf] + c1 = sCS[cs + 1] + s1 = sCS[cs + rhalf + 1] + r01 = _bf16_round( + _fma_rn(q0, c0, _mul_rn(_neg_f32(q1), s0)), + _fma_rn(q1, c0, _mul_rn(q0, s0)), + ) + r23 = _bf16_round( + _fma_rn(q2, c1, _mul_rn(_neg_f32(q3), s1)), + _fma_rn(q3, c1, _mul_rn(q2, s1)), + ) + else: + # NeoX: column d pairs with d +- rope / 2; four columns share a side. + lower = col < rhalf + pm4 = m4 - rhalf + cs = row * rope + col - rhalf + if lower: + pm4 = m4 + rhalf + cs = row * rope + col + p01 = _bf16_round(sAcc[pm4, row] * rstd, sAcc[pm4 + 1, row] * rstd) + p23 = _bf16_round(sAcc[pm4 + 2, row] * rstd, sAcc[pm4 + 3, row] * rstd) + p0 = p01[0] + p1 = p01[1] + p2 = p23[0] + p3 = p23[1] + if lower: + p0 = _neg_f32(p0) + p1 = _neg_f32(p1) + p2 = _neg_f32(p2) + p3 = _neg_f32(p3) + r01 = _bf16_round( + _fma_rn(q0, sCS[cs], _mul_rn(p0, sCS[cs + rhalf])), + _fma_rn(q1, sCS[cs + 1], _mul_rn(p1, sCS[cs + rhalf + 1])), + ) + r23 = _bf16_round( + _fma_rn(q2, sCS[cs + 2], _mul_rn(p2, sCS[cs + rhalf + 2])), + _fma_rn(q3, sCS[cs + 3], _mul_rn(p3, sCS[cs + rhalf + 3])), + ) + q0 = r01[0] + q1 = r01[1] + q2 = r23[0] + q3 = r23[1] + amax = cute.arch.fmax( + cute.arch.fmax(abs(q0), abs(q1)), cute.arch.fmax(abs(q2), abs(q3)) + ) + for i in cutlass.range_constexpr(int(math.log2(lanes))): + amax = cute.arch.fmax( + amax, cute.arch.shuffle_sync_bfly(amax, offset=lanes >> (i + 1)) + ) + rows.append(row) + m4s.append(m4) + vals.append((q0, q1, q2, q3)) + amaxes.append(amax) + + if cutlass.const_expr(self.pair): + # The other half's row amax: each owner sends its rows' amax to the other owner. + for g in cutlass.range_constexpr(groups_per_thread): + if m4s[g] == 0: + sAmaxMine[rows[g]] = amaxes[g] + cute.arch.barrier(barrier_id=_EPI_BARRIER_ID, number_of_threads=self.epilog_threads) + if epi_tid < self.cta_n // 4: + v = epi_tid * 4 + _store_shared_remote_v4( + sAmaxMine[v], + sAmaxMine[v + 1], + sAmaxMine[v + 2], + sAmaxMine[v + 3], + sAmaxPeer.iterator + v, + bar_amax, + Int32(1) - half, + ) + if epi_tid == 0: + cute.arch.mbarrier_arrive_and_expect_tx(bar_amax, self.cta_n * 4) + cute.arch.mbarrier_wait(bar_amax, 0) + for g in cutlass.range_constexpr(groups_per_thread): + amaxes[g] = cute.arch.fmax(amaxes[g], sAmaxPeer[rows[g]]) + + for g in cutlass.range_constexpr(groups_per_thread): + row = rows[g] + m4 = m4s[g] + q0, q1, q2, q3 = vals[g] + token = n_idx * self.cta_n + row + if token < n_tok: + scale = _fp8_ratio(amaxes[g]) + if cutlass.const_expr(qt.ue8m0): + scale = _ue8m0_up(scale) + inv = cute.arch.rcp_approx(scale) + feature = bidx * self.cta_m + m4 + addr = (mC.iterator + cute.crd2idx((feature, token, 0), mC.layout)).toint() + _st_global_b32( + addr, + _cvt_e4m3x4( + _mul_rn(q0, inv), _mul_rn(q1, inv), _mul_rn(q2, inv), _mul_rn(q3, inv) + ), + ) + if col0 + m4 == 0: + mQScaleOut[token, head] = scale + mWOut[token, head] = _mul_rn( + _mul_rn(sW[row], scale), Float32(qt.weight_scale) + ) + + +# ----------------------------------------------------------------------------------- host + + +@cute.experimental.jit +def _bmm_fused( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + gamma: cute.Tensor, + norm_out: cute.Tensor, + kv_lens: cute.Tensor, + cos_sin: cute.Tensor, + q_scale: cute.Tensor, + stream: _cuda.CUstream, +): + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])), + gamma, + norm_out, + kv_lens, + cos_sin, + q_scale, + stream, + ) + + +@functools.cache +def _placeholders(device_index: int) -> dict: + """Stand-ins for the operands an epilogue does not read (one compiled signature).""" + device = _torch.device("cuda", device_index) + return { + "gamma": _torch.zeros(_BF16_PER_VEC, dtype=_torch.bfloat16, device=device), + "norm_out": _torch.zeros((1, _BF16_PER_VEC), dtype=_torch.bfloat16, device=device), + "kv_lens": _torch.zeros(4, dtype=_torch.int32, device=device), + "cos_sin": _torch.zeros(4, dtype=_torch.float32, device=device), + "q_scale": _torch.ones(4, dtype=_torch.float32, device=device), + } + + +def _aux_cute_tensors(gamma, norm_out, kv_lens, cos_sin, q_scale): + # detach(): a module parameter (the norm weight) cannot be exported through DLPack otherwise. + return ( + _from_dlpack_dynamic(gamma.detach(), 0, assumed_align=16), + _from_dlpack_dynamic(norm_out.detach(), 1, assumed_align=16), + _from_dlpack_dynamic(kv_lens.detach(), 0, assumed_align=4), + _from_dlpack_dynamic(cos_sin.detach(), 0, assumed_align=4), + _from_dlpack_dynamic(q_scale.detach(), 0, assumed_align=4), + ) + + +def _swap_cute_tensors(a, b, out): + a_swap = b.detach().unsqueeze(0).transpose(-2, -1) + b_swap = a.detach().unsqueeze(0).transpose(-2, -1) + c_swap = out.detach().unsqueeze(0).transpose(-2, -1) + leading_dims = tuple(_detect_leading_dim(t) for t in (a_swap, b_swap, c_swap)) + tensors = tuple( + _from_dlpack_dynamic(t, ld) + for t, ld in zip((a_swap, b_swap, c_swap), leading_dims, strict=True) + ) + return tensors, leading_dims + + +@functools.cache +def _get_compiled_fused_kernel( + tactic: SplitKTactic, + use_pdl: bool, + epilogue: SplitKEpilogue, + leading_dims: tuple[int, int, int], + out_dtype: _torch.dtype, + reserve_full_smem: bool, + device_index: int, +): + kernel = SplitKFusedEpilogueGemmKernel( + tactic=tactic, use_pdl=use_pdl, epilogue=epilogue, reserve_full_smem=reserve_full_smem + ) + device = _torch.device("cuda", device_index) + m, n, k = 8, 64 * 9, _CTA_K * 2 + weight = _torch.empty((n, k), dtype=_torch.bfloat16, device=device) + act = _torch.empty((m, k), dtype=_torch.bfloat16, device=device) + out = _torch.empty((m, n), dtype=out_dtype, device=device) + tensors, repr_leading = _swap_cute_tensors(act, weight.t(), out) + if repr_leading != leading_dims: + raise ValueError(f"unsupported operand layouts {leading_dims} (expected {repr_leading})") + ph = _placeholders(device_index) + aux = _aux_cute_tensors( + ph["gamma"], ph["norm_out"], ph["kv_lens"], ph["cos_sin"], ph["q_scale"] + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(device).cuda_stream) + return cute_ext.compile(_bmm_fused, kernel, *tensors, *aux, stream) + + +def _check(cond: bool, msg: str) -> None: + if not cond: + raise ValueError(msg) + + +def _validate_fused(a, b, out, epilogue, gamma, norm_out, kv_lens, cos_sin, q_scale): + _check(all(isinstance(t, _torch.Tensor) for t in (a, b, out)), "a, b and out must be tensors") + _check(a.ndim == 2 and b.ndim == 2 and out.ndim == 2, "a, b and out must be 2D") + _check( + a.is_cuda and b.device == a.device and out.device == a.device, + "a, b and out must be on one CUDA device", + ) + _check(a.dtype == _torch.bfloat16 and b.dtype == _torch.bfloat16, "a and b must be BF16") + m, k = a.shape + _check(b.shape[0] == k, f"incompatible shapes {tuple(a.shape)} x {tuple(b.shape)}") + n = b.shape[1] + _check(tuple(out.shape) == (m, n), f"out must be {(m, n)}, got {tuple(out.shape)}") + _check( + 1 <= m <= 32 and k % _CTA_K == 0 and n % 4 == 0, f"unsupported shape M={m} N={n} K={k}" + ) + # a may be a row view of a wider buffer (the raw q_a slice of the fused-A output). + _check( + a.stride(1) == 1 and a.stride(0) >= k and (a.stride(0) * 2) % 16 == 0, + f"a needs unit inner stride and a 16-byte row pitch, got strides {a.stride()}", + ) + _check(b.t().is_contiguous(), "b must be column-major (the transpose of a contiguous [N, K])") + _check(out.is_contiguous(), "out must be contiguous") + _check( + all(t.data_ptr() % 32 == 0 for t in (a, b, out)), "a, b and out must be 32-byte aligned" + ) + if epilogue.fp8_rope: + _check(out.dtype == _torch.uint8, "fp8_rope writes uint8 (E4M3 bits)") + _check( + n % epilogue.head_dim == 0, f"N={n} is not a multiple of head_dim {epilogue.head_dim}" + ) + _check( + m % epilogue.seq_len == 0, f"M={m} is not a multiple of seq_len {epilogue.seq_len}" + ) + for t, name, dtype in ( + (kv_lens, "kv_lens", _torch.int32), + (cos_sin, "cos_sin", _torch.float32), + (q_scale, "q_scale", _torch.float32), + ): + _check( + isinstance(t, _torch.Tensor) + and t.dtype == dtype + and t.is_contiguous() + and t.device == a.device, + f"fp8_rope needs a contiguous {dtype} {name} on the GEMM's device", + ) + _check( + kv_lens.numel() >= m // epilogue.seq_len, "kv_lens has fewer entries than sequences" + ) + _check(q_scale.numel() >= 1, "q_scale is empty") + else: + _check(out.dtype == _torch.bfloat16, "out must be BF16 without fp8_rope") + if epilogue.row_rms: + _check(k == epilogue.rms_k, f"row_rms was specialized for K={epilogue.rms_k}, got {k}") + if epilogue.rms_out: + _check( + isinstance(gamma, _torch.Tensor) + and gamma.dtype == _torch.bfloat16 + and gamma.is_contiguous() + and gamma.numel() == k + and gamma.data_ptr() % 16 == 0, + "rms_out needs a contiguous 16-byte aligned BF16 gamma of K elements", + ) + _check( + isinstance(norm_out, _torch.Tensor) + and norm_out.dtype == _torch.bfloat16 + and norm_out.is_contiguous() + and tuple(norm_out.shape) == (m, k) + and norm_out.data_ptr() % 16 == 0, + "rms_out needs a contiguous 16-byte aligned BF16 norm_out [M, K]", + ) + return m, n, k + + +def run_splitk_fused( + a: _torch.Tensor, + b: _torch.Tensor, + out: _torch.Tensor, + pdl: bool, + tactic: SplitKTactic, + epilogue: SplitKEpilogue, + *, + gamma: _torch.Tensor | None = None, + norm_out: _torch.Tensor | None = None, + kv_lens: _torch.Tensor | None = None, + cos_sin: _torch.Tensor | None = None, + q_scale: _torch.Tensor | None = None, + reserve_full_smem: bool = True, +) -> _torch.Tensor: + """Run ``a[M, K] @ b[K, N]`` with the split-K kernel and a fused epilogue into ``out``. + + Args: + a: Activation [M, K] BF16 with unit inner stride (may be a row view of a wider + buffer). With ``epilogue.row_rms`` it is the raw RMSNorm input. + b: Weight [K, N], column-major (``W.t()`` of a contiguous [N, K]); with ``row_rms`` + the RMSNorm weight is already folded in. + out: [M, N] contiguous: BF16, or uint8 (E4M3 bits) with ``epilogue.fp8_rope``. + pdl: Launch with Programmatic Dependent Launch. + tactic: The split-K tactic; its A/B stages are reduced to fit the epilogue buffers. + epilogue: The fused epilogue. + gamma, norm_out: ``rms_out`` operands (BF16 [K], BF16 [M, K]). + kv_lens, cos_sin, q_scale: ``fp8_rope`` operands (int32 KV lengths per sequence, the + FP32 (cos, sin) table of ``mla_rope_generation``, FP32 [1] scale). + reserve_full_smem: Reserve the full shared-memory capacity per CTA. + """ + m, n, k = _validate_fused(a, b, out, epilogue, gamma, norm_out, kv_lens, cos_sin, q_scale) + tactic = fused_epilogue_tactic(tactic) + validate_tactic(tactic, m, n, k) + _check( + _fused_smem_bytes(tactic, tactic.ab_stages) <= _SMEM_CAPACITY_BYTES, + f"tactic {tactic} does not fit the epilogue buffers", + ) + tensors, leading_dims = _swap_cute_tensors(a, b, out) + device_index = a.device.index if a.device.index is not None else _torch.cuda.current_device() + compiled = _get_compiled_fused_kernel( + tactic, pdl, epilogue, leading_dims, out.dtype, reserve_full_smem, device_index + ) + ph = _placeholders(device_index) + aux = _aux_cute_tensors( + gamma if epilogue.rms_out else ph["gamma"], + norm_out if epilogue.rms_out else ph["norm_out"], + kv_lens if epilogue.fp8_rope else ph["kv_lens"], + cos_sin.view(-1) if epilogue.fp8_rope else ph["cos_sin"], + q_scale if epilogue.fp8_rope else ph["q_scale"], + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + compiled(*tensors, *aux, stream) + return out + + +@cute.experimental.jit +def _bmm_qtail( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + positions: cute.Tensor, + cos_sin: cute.Tensor, + weights_in: cute.Tensor, + q_scale_out: cute.Tensor, + weights_out: cute.Tensor, + stream: _cuda.CUstream, +): + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])), + positions, + cos_sin, + weights_in, + q_scale_out, + weights_out, + stream, + ) + + +def _qtail_cute_tensors(positions, cos_sin, weights_in, q_scale_out, weights_out): + return ( + _from_dlpack_dynamic(positions.detach(), 0, assumed_align=4), + _from_dlpack_dynamic(cos_sin.detach(), 1, assumed_align=4), + _from_dlpack_dynamic(weights_in.detach(), 1, assumed_align=2), + _from_dlpack_dynamic(q_scale_out.detach(), 1, assumed_align=4), + _from_dlpack_dynamic(weights_out.detach(), 1, assumed_align=4), + ) + + +@functools.cache +def _get_compiled_qtail_kernel( + tactic: SplitKTactic, + use_pdl: bool, + epilogue: SplitKEpilogue, + qtail: IndexerQTail, + leading_dims: tuple[int, int, int], + reserve_full_smem: bool, + device_index: int, +): + kernel = SplitKIndexerQTailGemmKernel( + tactic=tactic, + use_pdl=use_pdl, + epilogue=epilogue, + qtail=qtail, + reserve_full_smem=reserve_full_smem, + ) + device = _torch.device("cuda", device_index) + m, n, k = 8, qtail.n_heads * qtail.head_dim, epilogue.rms_k + weight = _torch.empty((n, k), dtype=_torch.bfloat16, device=device) + act = _torch.empty((m, k), dtype=_torch.bfloat16, device=device) + out = _torch.empty((m, n), dtype=_torch.uint8, device=device) + tensors, repr_leading = _swap_cute_tensors(act, weight.t(), out) + if repr_leading != leading_dims: + raise ValueError(f"unsupported operand layouts {leading_dims} (expected {repr_leading})") + aux = _qtail_cute_tensors( + _torch.zeros(m, dtype=_torch.int32, device=device), + _torch.zeros((m, qtail.rope_dim), dtype=_torch.float32, device=device), + _torch.zeros((m, qtail.n_heads), dtype=_torch.bfloat16, device=device), + _torch.zeros((m, qtail.n_heads), dtype=_torch.float32, device=device), + _torch.zeros((m, qtail.n_heads), dtype=_torch.float32, device=device), + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(device).cuda_stream) + return cute_ext.compile(_bmm_qtail, kernel, *tensors, *aux, stream) + + +def _validate_qtail( + a, b, out, epilogue, qtail, positions, cos_sin, weights_in, q_scale_out, weights_out +): + _check(all(isinstance(t, _torch.Tensor) for t in (a, b, out)), "a, b and out must be tensors") + _check(a.ndim == 2 and b.ndim == 2 and out.ndim == 2, "a, b and out must be 2D") + _check( + a.is_cuda and b.device == a.device and out.device == a.device, + "a, b and out must be on one CUDA device", + ) + _check(a.dtype == _torch.bfloat16 and b.dtype == _torch.bfloat16, "a and b must be BF16") + m, k = a.shape + _check(b.shape[0] == k, f"incompatible shapes {tuple(a.shape)} x {tuple(b.shape)}") + n = b.shape[1] + _check(tuple(out.shape) == (m, n), f"out must be {(m, n)}, got {tuple(out.shape)}") + _check(1 <= m <= 32 and k % _CTA_K == 0, f"unsupported shape M={m} N={n} K={k}") + _check( + a.stride(1) == 1 and a.stride(0) >= k and (a.stride(0) * 2) % 16 == 0, + f"a needs unit inner stride and a 16-byte row pitch, got strides {a.stride()}", + ) + _check(b.t().is_contiguous(), "b must be column-major (the transpose of a contiguous [N, K])") + _check(out.is_contiguous() and out.dtype == _torch.uint8, "out must be contiguous uint8") + _check( + all(t.data_ptr() % 32 == 0 for t in (a, b, out)), "a, b and out must be 32-byte aligned" + ) + _check( + epilogue.row_rms and not epilogue.rms_out and not epilogue.fp8_rope, + "the indexer Q tail follows the row_rms epilogue alone", + ) + _check(k == epilogue.rms_k, f"row_rms was specialized for K={epilogue.rms_k}, got {k}") + qtail.validate() + _check( + n == qtail.n_heads * qtail.head_dim, + f"N={n} is not {qtail.n_heads} heads of {qtail.head_dim}", + ) + _check( + isinstance(positions, _torch.Tensor) + and positions.dtype == _torch.int32 + and positions.dim() == 1 + and positions.is_contiguous() + and positions.numel() >= m, + "positions must be a contiguous int32 vector of at least M entries", + ) + _check( + isinstance(cos_sin, _torch.Tensor) + and cos_sin.dtype == _torch.float32 + and cos_sin.dim() == 2 + and cos_sin.stride(1) == 1 + and cos_sin.shape[1] >= qtail.rope_dim, + "cos_sin must be FP32 [positions, >= rope_dim] with unit inner stride", + ) + _check( + isinstance(weights_in, _torch.Tensor) + and weights_in.dtype == _torch.bfloat16 + and tuple(weights_in.shape) == (m, qtail.n_heads) + and weights_in.stride(1) == 1, + "weights_in must be BF16 [M, n_heads] with unit inner stride", + ) + for t, name in ((q_scale_out, "q_scale_out"), (weights_out, "weights_out")): + _check( + isinstance(t, _torch.Tensor) + and t.dtype == _torch.float32 + and tuple(t.shape) == (m, qtail.n_heads) + and t.is_contiguous(), + f"{name} must be a contiguous FP32 [M, n_heads]", + ) + for t in (positions, cos_sin, weights_in, q_scale_out, weights_out): + _check(t.device == a.device, "every operand must be on the GEMM's device") + return m, n, k + + +def run_splitk_idx_qtail( + a: _torch.Tensor, + b: _torch.Tensor, + out: _torch.Tensor, + pdl: bool, + tactic: SplitKTactic, + epilogue: SplitKEpilogue, + qtail: IndexerQTail, + *, + positions: _torch.Tensor, + cos_sin: _torch.Tensor, + weights_in: _torch.Tensor, + q_scale_out: _torch.Tensor, + weights_out: _torch.Tensor, + reserve_full_smem: bool = True, +) -> _torch.Tensor: + """The DSA indexer wq_b ``a[M, K] @ b[K, N]`` with the row-RMS epilogue and the indexer Q tail. + + Args: + a: The raw q_a [M, K] BF16 with unit inner stride (may be a row view of a wider buffer). + b: The norm-folded indexer wq_b [K, N], column-major (``W.t()`` of a contiguous [N, K]). + out: [M, N] uint8: the E4M3 indexer Q, [tokens, heads * head_dim]. + pdl: Launch with Programmatic Dependent Launch. + tactic: The split-K tactic: ``mma_m`` is one head or half a head; its A/B stages are + reduced to fit the epilogue buffers. + epilogue: ``row_rms`` alone. + qtail: The Q tail. + positions: int32 [>= M] RoPE positions of the tokens. + cos_sin: FP32 [positions, >= rope_dim] cos/sin rows ([cos | sin] halves). + weights_in: BF16 [M, n_heads] attention weights (the weights_proj output; may be a column + view of a wider buffer). + q_scale_out: FP32 [M, n_heads] output: the E4M3 scales. + weights_out: FP32 [M, n_heads] output: weights_in * scale * qtail.weight_scale. + reserve_full_smem: Reserve the full shared-memory capacity per CTA. + """ + m, n, k = _validate_qtail( + a, b, out, epilogue, qtail, positions, cos_sin, weights_in, q_scale_out, weights_out + ) + tactic = idx_qtail_tactic(tactic, qtail) + validate_tactic(tactic, m, n, k) + _check( + _qtail_smem_bytes(tactic, tactic.ab_stages, qtail) <= _SMEM_CAPACITY_BYTES, + f"tactic {tactic} does not fit the indexer Q tail buffers", + ) + tensors, leading_dims = _swap_cute_tensors(a, b, out) + device_index = a.device.index if a.device.index is not None else _torch.cuda.current_device() + compiled = _get_compiled_qtail_kernel( + tactic, pdl, epilogue, qtail, leading_dims, reserve_full_smem, device_index + ) + aux = _qtail_cute_tensors(positions, cos_sin, weights_in, q_scale_out, weights_out) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + compiled(*tensors, *aux, stream) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/draft_hidden_half.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/draft_hidden_half.py new file mode 100644 index 000000000000..99bdabd41f32 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/draft_hidden_half.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""This rank's FP32 share of the next MTP draft step's eh_proj hidden half, hnorm included, as +one split-K GEMM.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +from ._low_m_splitk import ( + device_limits, + operands_overlap, + run_splitk_hidden_half, + splitk_default_tactic, +) + +#: GLM-5.2 at tp16: hidden 6144, this rank's 6144 / 16 = 384 output rows of eh_proj's hidden +#: half, hnorm eps 1e-5. +_GLM = dict(k=6144, n=384, eps=1e-5) + + +class GlmDraftHiddenHalf(OpWrapper): + """`out = rmsnorm(x) @ weight.T` (with `norm_weight`) or `out = x @ weight.T` (without), FP32 + out; returns `out`. + + eh_proj is row-parallel over `cat(enorm(embed(token)), hnorm(hidden))`, so its output is a + token half (tabulated at load) plus `hnorm(hidden) @ W_e[:, H:].T`. This op computes this + rank's `H / tp` outputs of that hidden half in FP32, on the draft step's hidden state before + the draft LM head runs. `weight` is this rank's rows `[rank * H / tp, (rank + 1) * H / tp)` + of `W_e[:, H:]` (gathered from the row-parallel shards once, at load). + + With `norm_weight` (one row of H = 6144): the kernel's B producer computes FlashInfer's + RMSNorm of `x` itself -- the 64-thread layout, the sequential square sums, the two + butterflies and their sum, `rsqrt.approx.ftz` of the IEEE `sum / H + eps`, `(x * rstd) * + (w + 0)` rounded to BF16 -- so the result is bit-identical to running that RMSNorm first and + the GEMM on its output. Without it `x` is the normalized row(s). + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="MTP draft steps 0-5 (one request): 1 row, hnorm in the kernel, split-K 8 " + "(48 CTAs, every K tile in its own stage)", + spec=dict(_GLM, rows=1, in_kernel_norm=True), + ), + Cell( + why="two requests per draft step: 2 normalized rows, the plain FP32-output GEMM", + spec=dict(_GLM, rows=2, in_kernel_norm=False), + ), + ) + + note = """ + The tactic is the shape's `splitk_default_tactic`; the in-kernel norm needs one row, H = + 6144 and a tactic whose CTA K slice fits its stages (`hnorm_b_supported`), else ValueError. + PDL follows `TRTLLM_ENABLE_PDL` (default on): before its grid dependency wait the kernel + reads only the weight (TMA) and the norm weight and zeroes its own B stages; it reads `x` + after the wait. + + Loud, so not guarded (ValueError before the launch): M outside [1, 32], K not a multiple of + 128, non-BF16 x or weight, a non-FP32 or non-contiguous out, pointers not 32-byte aligned, + the in-kernel norm where it does not serve. The first call of a (tactic, norm eps) compiles + the kernel; make it before CUDA-graph capture. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + norm_weight: torch.Tensor | None, + eps: float | None, + ) -> torch.Tensor: + """Write `out` and return it. + + Args: + x: [rows, H] BF16 with unit inner stride: the raw hidden state with `norm_weight`, + else the normalized one. + weight: [N, H] BF16, contiguous: this rank's hidden-half rows of eh_proj. + out: [rows, N] FP32, contiguous. + norm_weight: [H] BF16 hnorm weight for the in-kernel norm, or None. + eps: hnorm's epsilon (with `norm_weight`), else None. + """ + rows, k = x.shape + sm_count, smem_capacity = device_limits(x.device.index) + tactic = splitk_default_tactic( + rows, weight.shape[0], k, sm_count=sm_count, smem_capacity=smem_capacity + ) + return run_splitk_hidden_half( + x, weight, out, get_env_enable_pdl(), tactic, norm_weight=norm_weight, norm_eps=eps + ) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + norm_weight: torch.Tensor | None, + eps: float | None, + ) -> torch.Tensor: + """FP32 product (TF32 off) of the BF16 normalized row(s) and the weight. The norm is + FlashInfer's formula in FP32 -- `rstd = rsqrt(sum(x ** 2) / H + eps)`, `(x * rstd) * w` + rounded to BF16 once -- with torch's summation order. `out` is not read.""" + xf = x.float() + if norm_weight is not None: + rstd = torch.rsqrt(xf.pow(2).sum(dim=-1, keepdim=True) / xf.shape[-1] + eps) + xf = ((xf * rstd) * norm_weight.float()).to(torch.bfloat16).float() + with true_fp32_matmul(): + return xf @ weight.float().t() + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + norm_weight: torch.Tensor | None, + eps: float | None, + ) -> None: + # `eps` without `norm_weight` is ignored: the row goes into the GEMM unnormalized. + assert (norm_weight is None) == (eps is None), ( + "pass norm_weight and eps together: eps alone leaves x unnormalized" + ) + # The owner CTA of a tile stores while other clusters still read x and the weight. + for operand, name in ((x, "x"), (weight, "weight")): + assert not operands_overlap(out, operand), f"out must not share memory with {name}" + + +draft_hidden_half = GlmDraftHiddenHalf() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/indexer_q_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/indexer_q_tail.py new file mode 100644 index 000000000000..98dcc783122e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/indexer_q_tail.py @@ -0,0 +1,253 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DSA indexer wq_b from the raw q_a as one split-K GEMM whose epilogue finishes the indexer Q: +RoPE, the per-head E4M3 scale, the E4M3 Q and the scaled head weights.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +from ._low_m_splitk import operands_overlap +from ._low_m_splitk_fused import ( + IndexerQTail, + SplitKEpilogue, + idx_qtail_plan, + run_splitk_idx_qtail, +) + +#: GLM-5.2's DSA indexer (replicated on every rank): 32 heads of 128, RoPE on the first 64 +#: columns as GPT-J pairs, UE8M0 scales, weight scale 128 ** -0.5 * 32 ** -0.5; q_lora_rank 2048 +#: and the q_a RMSNorm eps 1e-5. +_GLM = dict( + n_heads=32, + head_dim=128, + rope_dim=64, + interleave=True, + ue8m0=True, + weight_scale=128**-0.5 * 32**-0.5, + k=2048, + eps=1e-5, +) + +#: The E4M3 ratio of the indexer prologue: max(amax, 1e-12) * fp32(1 / 448). +_AMAX_FLOOR = 1e-12 +_INV_E4M3_MAX = 1.0 / 448.0 + + +class GlmIndexerQTail(OpWrapper): + """The DSA indexer's FP8 Q, Q scales and scaled head weights from the raw q_a; returns + `(q_fp8, q_scale, weights_out)`. + + The GEMM is the indexer wq_b with the q_a RMSNorm folded into its weight (`weight` is + `W_wq_b * gamma_q_a`, per column, built in FP64 and rounded to BF16 once at load): `y = bf16( + rstd_t * (q_a @ weight.T))`, `rstd_t = rsqrt(mean(q_a[t] ** 2) + eps)`, rounded to BF16 as a + BF16 store would. Then, per (token t, head h) of `y`, the indexer prologue's tail in its + order: RoPE on the head's first `rope_dim` columns (GPT-J pairs with `interleave`, NeoX + halves otherwise; cos/sin from the row of `positions[t]`) as `bf16(fma(x, cos, rn(+-partner * + sin)))`; `scale = max(amax, 1e-12) * fp32(1 / 448)`, rounded up to a power of two with + `ue8m0`; `q_fp8 = e4m3(q * (1 / scale))`; `q_scale = scale`; `weights_out = (weights_in * + scale) * weight_scale`. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="full-indexer layer, verify step and MTP draft step 0: 8 tokens, 32 heads of 128; " + "half-head tiles in (2, 2) clusters (128 CTAs)", + spec=dict(_GLM, m=8), + ), + Cell( + why="1 token, decode without MTP: one row in an 8-row tile (the same tactic)", + spec=dict(_GLM, m=1), + ), + ) + + note = """ + Kernel: `SplitKIndexerQTailGemmKernel`, the row-RMS fused split-K GEMM at the shape's + `splitk_default_tactic` (64-row tiles, split-K 2 for 32 x 128 outputs from 2048) whose + 64-row tiles hold half a head; the two halves of a head run in one cluster whose owner CTAs + exchange their row amax through DSMEM. The owner's epilogue warps compute the row RMS after + the grid dependency wait, then read the positions, the cos/sin rows and weights_in while the + mainloop runs. PDL follows `TRTLLM_ENABLE_PDL` (default on): before its wait the kernel reads + only the weight. + + The values are those of the fused indexer prologue run on wq_b's BF16 output: the same + operations in the same order. The interleaved RoPE and UE8M0 scales are GLM-5.2's; the NeoX + form and plain FP32 scales are implemented but not cells. + + Loud, so not guarded (ValueError before the launch): M outside [1, 32], K not a multiple of + 128, N not n_heads * head_dim, a head_dim that is neither one nor two 64-row tiles, operands + of the wrong dtype, shape or stride, a rope_dim the epilogue cannot split. The first call of + a (tactic, Q tail) compiles the kernel; make it before CUDA-graph capture. + """ + + def __call__( + self, + q_a: torch.Tensor, + weight: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + weights_in: torch.Tensor, + q_fp8: torch.Tensor, + q_scale: torch.Tensor, + weights_out: torch.Tensor, + eps: float, + n_heads: int, + head_dim: int, + rope_dim: int, + interleave: bool, + ue8m0: bool, + weight_scale: float, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Write `q_fp8`, `q_scale` and `weights_out` and return them. + + Args: + q_a: [M, q_lora_rank] BF16, the raw q_a (before its RMSNorm); unit inner stride, a + row pitch of a multiple of 8 elements (a row view of the fused-A output works). + weight: [n_heads * head_dim, q_lora_rank] BF16, contiguous: wq_b, norm-folded. + positions: [>= M] int32 RoPE positions, contiguous. + cos_sin: [table positions, >= rope_dim] FP32 with unit inner stride: per position + the rope_dim / 2 cosines, then the rope_dim / 2 sines. + weights_in: [M, n_heads] BF16 head weights (weights_proj output) with unit inner + stride (a column view of the fused-A output works). + q_fp8: [M, n_heads, head_dim] float8_e4m3fn, contiguous (output). + q_scale: [M, n_heads, 1] FP32, contiguous (output). + weights_out: [M, n_heads] FP32, contiguous (output). + eps: The q_a RMSNorm epsilon. + n_heads, head_dim, rope_dim, interleave, ue8m0, weight_scale: The indexer's Q tail. + """ + m, k = q_a.shape + qtail = IndexerQTail( + n_heads=n_heads, + head_dim=head_dim, + rope_dim=rope_dim, + interleave=interleave, + ue8m0=ue8m0, + weight_scale=float(weight_scale), + ) + tactic = idx_qtail_plan(qtail, m, weight.shape[0], k, q_a.device.index) + run_splitk_idx_qtail( + q_a, + weight.t(), + q_fp8.view(torch.uint8).view(m, -1), + get_env_enable_pdl(), + tactic, + SplitKEpilogue(row_rms=True, rms_k=k, rms_eps=float(eps)), + qtail, + positions=positions, + cos_sin=cos_sin, + weights_in=weights_in, + q_scale_out=q_scale.view(m, n_heads), + weights_out=weights_out, + ) + return q_fp8, q_scale, weights_out + + def reference( + self, + q_a: torch.Tensor, + weight: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + weights_in: torch.Tensor, + q_fp8: torch.Tensor, + q_scale: torch.Tensor, + weights_out: torch.Tensor, + eps: float, + n_heads: int, + head_dim: int, + rope_dim: int, + interleave: bool, + ue8m0: bool, + weight_scale: float, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """`(q_fp8, q_scale, weights_out)` as new tensors. The GEMM in FP32 (TF32 off); every + rounding the prologue applies, applied here; each RoPE fma in FP64, rounded to FP32 + once. The output arguments are not read.""" + m = q_a.shape[0] + x = q_a.float() + rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) + with true_fp32_matmul(): + y = ((x @ weight.float().t()) * rstd).to(torch.bfloat16).float() + y = y.view(m, n_heads, head_dim) + half = rope_dim // 2 + rows = cos_sin.reshape(cos_sin.shape[0], -1)[positions[:m].long(), :rope_dim] + cos = rows[:, :half].unsqueeze(1) + sin = rows[:, half:rope_dim].unsqueeze(1) + pe = y[..., :rope_dim] + if interleave: + first, second = pe[..., 0::2], pe[..., 1::2] + else: + first, second = pe[..., :half], pe[..., half:] + + def fma(a, b, c): + return (a.double() * b.double() + c.double()).float() + + rot_first = fma(first, cos, -second * sin).to(torch.bfloat16).float() + rot_second = fma(second, cos, first * sin).to(torch.bfloat16).float() + if interleave: + pe = torch.stack((rot_first, rot_second), dim=-1).flatten(-2) + else: + pe = torch.cat((rot_first, rot_second), dim=-1) + q = torch.cat((pe, y[..., rope_dim:]), dim=-1) + ratio = q.abs().amax(dim=-1, keepdim=True).clamp_min(_AMAX_FLOOR) * torch.tensor( + _INV_E4M3_MAX, dtype=torch.float32 + ) + scale = _ue8m0_round_up(ratio) if ue8m0 else ratio + q_ref = (q * (1.0 / scale)).clamp(-448.0, 448.0).to(torch.float8_e4m3fn) + weights_ref = (weights_in.float() * scale.squeeze(-1)) * torch.tensor( + weight_scale, dtype=torch.float32 + ) + return q_ref, scale, weights_ref + + def is_valid( + self, + q_a: torch.Tensor, + weight: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + weights_in: torch.Tensor, + q_fp8: torch.Tensor, + q_scale: torch.Tensor, + weights_out: torch.Tensor, + eps: float, + n_heads: int, + head_dim: int, + rope_dim: int, + interleave: bool, + ue8m0: bool, + weight_scale: float, + ) -> None: + # A position outside the cos/sin table reads past it: garbage, not an error. + used = positions[: q_a.shape[0]].long() + assert bool((used >= 0).all()) and int(used.max()) < cos_sin.shape[0], ( + "every position must index a row of the cos/sin table" + ) + # The epilogue stores while other clusters still read q_a, the weight and weights_in. + outputs = {"q_fp8": q_fp8, "q_scale": q_scale, "weights_out": weights_out} + operands = {"q_a": q_a, "weight": weight, "weights_in": weights_in} + for output_name, output in outputs.items(): + for name, operand in operands.items(): + assert not operands_overlap(output, operand), ( + f"{output_name} must not share memory with {name}" + ) + + +def _ue8m0_round_up(ratio: torch.Tensor) -> torch.Tensor: + """The positive FP32 `ratio` with its mantissa cleared and, if it was nonzero, its exponent + raised by one.""" + bits = ratio.view(torch.int32) + exponent = bits & 0x7F800000 + exponent = exponent + ((bits & 0x007FFFFF) != 0).to(torch.int32) * 0x00800000 + return exponent.view(torch.float32) + + +indexer_q_tail = GlmIndexerQTail() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/low_m_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/low_m_gemm.py new file mode 100644 index 000000000000..048993cf89c8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/low_m_gemm.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-M BF16 projection on the target's split-K CuTe DSL GEMM, into a caller-owned output.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +from ._low_m_splitk import ( + device_limits, + operands_overlap, + run_splitk_dense, + splitk_default_tactic, +) + + +class GlmLowMGemm(OpWrapper): + """`out = x @ weight.T`, BF16 in and out, FP32 accumulation rounded once; returns `out`. + + One launch of the low-M split-K kernel at the shape's `splitk_default_tactic`: each CTA of a + cluster accumulates an exact K slice in FP32, the cluster reduces through DSMEM in a fixed + order, and the owner CTA rounds and stores. The decode step's plain BF16 projections of at + most 32 rows take it: the shared expert's gate_up and down, the LM head on this rank's + vocabulary shard, and MTP draft step 0's eh_proj on this rank's input slice. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # GLM-5.2 at tp16: hidden 6144, shared-expert intermediate 2048 / 16 = 128 (gate_up fused, + # 256 rows), vocabulary 154880 / 16 = 9680, eh_proj input 2 * 6144 / 16 = 768. A verify + # step runs 8 rows (MTP K=7); the MTP draft steps run the shared expert and the LM head on + # 1 row, draft step 0's eh_proj on 8. + CELLS: tuple[Cell, ...] = ( + Cell( + why="shared expert gate_up, verify step: 8 x 6144 -> 256, split-K 8 (32 CTAs)", + spec=dict(m=8, k=6144, n=256), + ), + Cell( + why="shared expert gate_up, MTP draft step: 1 row", + spec=dict(m=1, k=6144, n=256), + ), + Cell( + why="shared expert down, verify step: 8 x 128 -> 6144, one K tile, no split (96 CTAs)", + spec=dict(m=8, k=128, n=6144), + ), + Cell( + why="shared expert down, MTP draft step: 1 row", + spec=dict(m=1, k=128, n=6144), + ), + Cell( + why="LM head vocabulary shard, verify step: 8 x 6144 -> 9680, partial last tile", + spec=dict(m=8, k=6144, n=9680), + ), + Cell( + why="LM head vocabulary shard, MTP draft step: 1 row", + spec=dict(m=1, k=6144, n=9680), + ), + Cell( + why="eh_proj, MTP draft step 0: 8 x 768 (row-parallel slice) -> 6144, split-K 2", + spec=dict(m=8, k=768, n=6144), + ), + ) + + note = """ + The tactic is a function of (M, N, K) and the device's SM count and opt-in shared memory + (`splitk_default_tactic`); the launch reserves the full shared-memory capacity, one CTA per + SM. PDL follows `TRTLLM_ENABLE_PDL` (default on): the weight is streamed before the grid + dependency wait, `x` is read after it, and the kernel releases its dependents once its weight + loads are issued. + + Loud, so not guarded: M outside [1, 32], K not a multiple of 128, mismatched shapes, dtypes + or devices, operands that are neither row- nor column-major, pointers not 32-byte aligned + (ValueError before the launch). The first call of a (shape class, tactic) compiles the + kernel; make it before CUDA-graph capture. + """ + + def __call__(self, x: torch.Tensor, weight: torch.Tensor, out: torch.Tensor) -> torch.Tensor: + """x [M, K] BF16 row-major; weight [N, K] BF16 (the stored weight); out [M, N] BF16.""" + m, k = x.shape + n = weight.shape[0] + sm_count, smem_capacity = device_limits(x.device.index) + tactic = splitk_default_tactic(m, n, k, sm_count=sm_count, smem_capacity=smem_capacity) + return run_splitk_dense(x, weight.t(), None, out, get_env_enable_pdl(), tactic) + + def reference(self, x: torch.Tensor, weight: torch.Tensor, out: torch.Tensor) -> torch.Tensor: + """The product the op writes into `out`: FP32-accumulated (TF32 off), rounded once. + + `out` is named to mirror `__call__` and not read. + """ + with true_fp32_matmul(): + return (x.float() @ weight.float().t()).to(x.dtype) + + def is_valid(self, x: torch.Tensor, weight: torch.Tensor, out: torch.Tensor) -> None: + # The owner CTA of a tile stores its rows while other clusters still stream x and the + # weight; an `out` that shares memory with either is overwritten under their reads and + # the result is silently wrong. + for operand, name in ((x, "x"), (weight, "weight")): + assert not operands_overlap(out, operand), f"out must not share memory with {name}" + + +low_m_gemm = GlmLowMGemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/o_proj_publish.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/o_proj_publish.py new file mode 100644 index 000000000000..af36c46b03b8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/o_proj_publish.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Attention o_proj whose split-K epilogue stores the BF16 output straight into this rank's slot of +the post-attention MNNVL one-shot all-reduce's current Lamport buffer.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +from ._low_m_splitk import ( + LamportPublish, + device_limits, + operands_overlap, + publish_tactic, + run_splitk_publish, + splitk_default_tactic, +) + +#: GLM-5.2 o_proj at tp16: 4 heads * v_head_dim 256 = 1024 -> hidden 6144, 16 ranks. +_GLM = dict(k=1024, n=6144, world_size=16) + +#: Lamport buffers of an MNNVL one-shot workspace, and the int32 words of its buffer flags. +_LAMPORT_BUFFERS = 3 +_FLAG_WORDS = 4 + + +class GlmOProjPublish(OpWrapper): + """`x @ weight.T` rounded to BF16, stored at `[token][rank][column]` of the all-reduce + workspace's current Lamport buffer instead of into an output; returns None. + + The current buffer is `lamport` (bytes) + `flags[0] * flags[2]`: the flags' current index + times their bytes per buffer, read once per CTA after the grid dependency wait (the previous + all-reduce rotated them). Value (t, c) lands at element `(t * world_size + rank) * N + c` of + that buffer, in BF16 with -0.0 written as +0.0 (-0.0 is the buffer's idle value): this + rank's shard as the one-shot kernel would have stored it. The all-reduce of this output must + then take the published shard (it neither loads nor stores its own) and poll every slot. The + owner CTA of each 64 x 8 tile stores 16-byte vectors of 8 columns through `lamport`, which is + the workspace's multicast address in production, so the slot fills on every rank. + + `out` gives the GEMM its output layout and is not written. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="verify step: 8 tokens, 1024 -> 6144, split-K 2 (192 CTAs), rank 5 of 16", + spec=dict(_GLM, m=8, rank=5, current=1), + ), + Cell( + why="MTP draft steps 1-6: 1 token, last rank, buffer index 2", + spec=dict(_GLM, m=1, rank=15, current=2), + ), + Cell( + why="verify step at 4 ranks: the slot arithmetic of another group size, buffer 0", + spec=dict(_GLM, m=8, rank=2, current=0, world_size=4), + ), + ) + + note = """ + The tactic is the shape's `splitk_default_tactic` with its A/B stages lowered to fit the FP32 + staging tile (`publish_tactic`; the c1 o_proj keeps all 12). PDL: the kernel is always + expected to run with `TRTLLM_ENABLE_PDL` on (the default) in front of the all-reduce; its + epilogue warps wait on the grid dependency, release the dependents (the all-reduce) right + after the wait and only then read the flags. With PDL off the flags are read without a wait, + which is ordered by the plain stream dependency. + + The Lamport workspace, its multicast mapping and its flags belong to the all-reduce; this op + reads the flags and writes one slot per token, nothing else. Loud, so not guarded + (ValueError before the launch): M outside [1, 32], K not a multiple of 128, N not a multiple + of 8, non-BF16 operands, a `lamport` that is not a 1-D BF16 view or `flags` that is not a + 1-D int32 view. The first call of a (tactic, world size, rank) compiles the kernel; make it + before CUDA-graph capture. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + lamport: torch.Tensor, + flags: torch.Tensor, + world_size: int, + rank: int, + ) -> None: + """Publish `x @ weight.T` into the current Lamport buffer. + + Args: + x: [M, K] BF16, contiguous: the attention output of this rank's heads. + weight: [N, K] BF16, contiguous: o_proj's row-parallel shard. + out: [M, N] BF16, contiguous, 32-byte aligned: layout only, not written. + lamport: 1-D BF16 view of the workspace's Lamport buffers (multicast address). + flags: 1-D int32 view of the workspace's buffer flags, 16-byte aligned: [0] the + current buffer index, [2] the bytes per buffer. + world_size: Ranks of the all-reduce group. + rank: This rank in the group. + """ + m, k = x.shape + n = weight.shape[0] + sm_count, smem_capacity = device_limits(x.device.index) + tactic = publish_tactic( + splitk_default_tactic(m, n, k, sm_count=sm_count, smem_capacity=smem_capacity), + smem_capacity=smem_capacity, + ) + run_splitk_publish( + x, + weight.t(), + out, + lamport, + flags, + get_env_enable_pdl(), + tactic, + LamportPublish(world_size=world_size, rank=rank, trigger_after_wait=True), + ) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + lamport: torch.Tensor, + flags: torch.Tensor, + world_size: int, + rank: int, + ) -> torch.Tensor: + """The [M, N] BF16 values the op stores at slots `t * world_size + rank` of the current + buffer: FP32-accumulated (TF32 off), rounded once, -0.0 as +0.0. Where they land is the + test's to check; `out`, `lamport` and `flags` are not read.""" + with true_fp32_matmul(): + values = (x.float() @ weight.float().t()).to(torch.bfloat16) + return torch.where(values == 0, torch.zeros_like(values), values) + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + lamport: torch.Tensor, + flags: torch.Tensor, + world_size: int, + rank: int, + ) -> None: + # A rank outside the group lands in another token's slots. + assert 0 <= rank < world_size, f"rank {rank} is outside a group of {world_size}" + # Every store goes through the flags' current buffer; slots past its end land in the + # next buffer (or past the workspace) without complaint. + assert flags.numel() >= _FLAG_WORDS and flags.data_ptr() % 16 == 0, ( + "flags must hold the four 16-byte aligned words the kernel reads" + ) + current, _, bytes_per_buffer, _ = (int(v) for v in flags[:_FLAG_WORDS].tolist()) + m, n = x.shape[0], weight.shape[0] + assert 0 <= current < _LAMPORT_BUFFERS, f"current buffer index {current} is out of range" + assert m * world_size * n * lamport.element_size() <= bytes_per_buffer, ( + "the Lamport buffer is smaller than the [token][rank][column] slots it must hold" + ) + assert (current + 1) * bytes_per_buffer <= lamport.numel() * lamport.element_size(), ( + "the current Lamport buffer runs past the end of `lamport`" + ) + # The epilogue stores while other clusters still read x and the weight. + for operand, name in ((x, "x"), (weight, "weight")): + assert not operands_overlap(lamport, operand), ( + f"lamport must not share memory with {name}" + ) + + +o_proj_publish = GlmOProjPublish() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/q_b_fp8_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/q_b_fp8_rope.py new file mode 100644 index 000000000000..44637c3ea359 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/q_b_fp8_rope.py @@ -0,0 +1,213 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MLA q_b of a decode step as one split-K GEMM: k_b absorbed, the q_a RMSNorm folded, and the +RoPE'd E4M3 attention Q written by the epilogue.""" + +from __future__ import annotations + +from typing import Literal + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +from ._low_m_splitk import operands_overlap +from ._low_m_splitk_fused import SplitKEpilogue, q_b_plan, run_splitk_fused + +#: GLM-5.2 MLA at tp16: 4 heads per rank, kv_lora_rank 512 + qk_rope_head_dim 64 per head, +#: q_lora_rank 2048, RMSNorm eps 1e-5. +_GLM = dict(heads=4, head_dim=576, rope_dim=64, k=2048, eps=1e-5) + +#: E4M3's largest finite magnitude. +_E4M3_MAX = 448.0 + + +class GlmQBFp8Rope(OpWrapper): + """The generation attention Q of a GLM MLA layer from the raw q_a, as E4M3 bits; returns `out`. + + `out[t, h * head_dim + c] = e4m3(q_scale * rope_t(rstd_t * (q_a @ weight.T))[t, h, c])`, where + `rstd_t = rsqrt(mean(q_a[t] ** 2) + eps)` and `rope_t` rotates the last `rope_dim` columns of + every `head_dim`-wide head as GPT-J pairs (2p, 2p + 1) by the (cos, sin) pair p of the token's + position, `kv_lens[s] - seq_len + i` for token i of sequence s (`seq_len` tokens each). The + conversion rounds to nearest even and saturates at +-448. Everything is FP32 from the GEMM's + accumulator to the conversion; nothing is rounded to BF16 on the way. + + `weight` is q_b with k_b absorbed and the q_a norm weight folded in, built once at load (see + the target's weights): for head h, rows [h * head_dim, h * head_dim + kv_lora_rank) are + `k_b_trans[h] @ W_q_b[h, :qk_nope]` (`k_b_trans` the MLA's [heads, kv_lora_rank, qk_nope] + k_b) and the next `rope_dim` rows are `W_q_b[h, qk_nope:]`, every column k times + `gamma_q_a[k]`, computed in FP64 and rounded to BF16 once. The output is the + `[tokens, heads, kv_lora_rank + rope_dim]` FP8 Q the sparse-MLA generation kernel reads. + + `plan` picks the launch, never the result's definition: `regular` (the shape's + `splitk_default_tactic`), `beside_topk` (128-row tiles, no split-K, three stages, for the + copy queued beside the DSA indexer top-k), `early_gather` (split-K 4 where the regular tactic + is split-K 8, beside a flash decode that gathers before its wait). The plans differ in the + FP32 summation order only. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="full-indexer layer, verify step and MTP draft step 0: 8 tokens of one request, " + "beside the top-k (128-row tiles, no split-K, 18 CTAs)", + spec=dict(_GLM, m=8, seq_len=8, plan="beside_topk"), + ), + Cell( + why="shared-indexer layer, verify step: 8 tokens, early-gather split-K 4 (72 CTAs)", + spec=dict(_GLM, m=8, seq_len=8, plan="early_gather"), + ), + Cell( + why="MTP draft steps 1-6: 1 token, regular split-K 8 (144 CTAs)", + spec=dict(_GLM, m=1, seq_len=1, plan="regular"), + ), + Cell( + why="two requests of 8 tokens: per-sequence positions, 16-row tiles (regular plan)", + spec=dict(_GLM, m=16, seq_len=8, plan="regular"), + ), + ) + + note = """ + Kernel: the split-K GEMM's mainloop and DSMEM reduction with the owner CTA's epilogue staged + in shared memory. The owner's epilogue warps compute the row RMS from global memory while the + mainloop runs (after the grid dependency wait) and read kv_lens, the cos/sin rows and + q_scale after the MMA. PDL follows `TRTLLM_ENABLE_PDL` (default on): before its wait the + kernel reads only the weight. + + Loud, so not guarded (ValueError before the launch): M outside [1, 32], K not a multiple of + 128, N not a multiple of head_dim, M not a multiple of seq_len, too few kv_lens, an `out` + that is not contiguous uint8, a weight that is not a contiguous [N, K], a q_a without unit + inner stride and a 16-byte row pitch, a geometry the epilogue rejects. The first call of a + (tactic, epilogue) compiles the kernel; make it before CUDA-graph capture. + """ + + def __call__( + self, + q_a: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kv_lens: torch.Tensor, + cos_sin: torch.Tensor, + q_scale: torch.Tensor, + eps: float, + head_dim: int, + rope_dim: int, + seq_len: int, + plan: Literal["regular", "beside_topk", "early_gather"], + ) -> torch.Tensor: + """Write `out` and return it. + + Args: + q_a: [M, q_lora_rank] BF16, the raw q_a (before its RMSNorm); unit inner stride, a + row pitch of a multiple of 8 elements (a row view of the fused-A output works). + weight: [heads * head_dim, q_lora_rank] BF16, contiguous: absorbed and norm-folded. + out: [M, heads * head_dim] uint8, contiguous: E4M3 bits. + kv_lens: [sequences] int32 KV lengths including this step's tokens. + cos_sin: [positions, rope_dim, 2] FP32 contiguous (cos, sin) pairs; position p's + pair i rotates columns 2i and 2i + 1 of every head's RoPE slice (the first + rope_dim / 2 pairs of a row are read). + q_scale: [1] FP32, the factor applied before the E4M3 conversion. + eps: The q_a RMSNorm epsilon. + head_dim: kv_lora_rank + rope_dim. + rope_dim: The RoPE width at the end of every head. + seq_len: Tokens per sequence; M is sequences * seq_len. + plan: The launch plan. + """ + m, k = q_a.shape + tactic = q_b_plan(plan, m, weight.shape[0], k, q_a.device.index) + epilogue = SplitKEpilogue( + row_rms=True, + rms_k=k, + rms_eps=float(eps), + fp8_rope=True, + head_dim=head_dim, + rope_dim=rope_dim, + seq_len=seq_len, + ) + return run_splitk_fused( + q_a, + weight.t(), + out, + get_env_enable_pdl(), + tactic, + epilogue, + kv_lens=kv_lens, + cos_sin=cos_sin, + q_scale=q_scale, + ) + + def reference( + self, + q_a: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kv_lens: torch.Tensor, + cos_sin: torch.Tensor, + q_scale: torch.Tensor, + eps: float, + head_dim: int, + rope_dim: int, + seq_len: int, + plan: Literal["regular", "beside_topk", "early_gather"], + ) -> torch.Tensor: + """The E4M3 bits the op writes, as uint8 [M, N]: FP32 throughout (TF32 off), one + conversion at the end. `out` and `plan` are not read.""" + m = q_a.shape[0] + x = q_a.float() + rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) + with true_fp32_matmul(): + q = (x @ weight.float().t()) * rstd + q = q.view(m, -1, head_dim) + nope = head_dim - rope_dim + token = torch.arange(m, device=q.device) + seq = token // seq_len + position = kv_lens.long()[seq] - seq_len + (token - seq * seq_len) + pairs = cos_sin.view(-1, rope_dim, 2)[position, : rope_dim // 2] + cos = pairs[..., 0].unsqueeze(1) + sin = pairs[..., 1].unsqueeze(1) + pe = q[..., nope:].reshape(m, q.shape[1], rope_dim // 2, 2) + even, odd = pe[..., 0], pe[..., 1] + rotated = torch.stack((cos * even - sin * odd, cos * odd + sin * even), dim=-1) + q = torch.cat((q[..., :nope], rotated.flatten(-2)), dim=-1) * q_scale.float()[0] + q = q.clamp(-_E4M3_MAX, _E4M3_MAX).to(torch.float8_e4m3fn) + return q.view(torch.uint8).reshape(m, -1) + + def is_valid( + self, + q_a: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kv_lens: torch.Tensor, + cos_sin: torch.Tensor, + q_scale: torch.Tensor, + eps: float, + head_dim: int, + rope_dim: int, + seq_len: int, + plan: Literal["regular", "beside_topk", "early_gather"], + ) -> None: + # The kernel indexes the flat table as (position * rope_dim + pair) * 2: a table of + # another row width is read at the wrong rows without complaint. + assert cos_sin.dim() == 3 and tuple(cos_sin.shape[1:]) == (rope_dim, 2), ( + f"cos_sin must be [positions, rope_dim={rope_dim}, 2], got {tuple(cos_sin.shape)}" + ) + # A position outside the table reads past it (or before it): garbage, not an error. + sequences = q_a.shape[0] // seq_len + lens = kv_lens[:sequences].long() + assert bool((lens >= seq_len).all()), "kv_lens must cover the step's own tokens" + assert int(lens.max()) <= cos_sin.shape[0], ( + "the cos/sin table must hold every position up to max(kv_lens) - 1" + ) + # The epilogue stores while other clusters still read q_a and the weight. + for operand, name in ((q_a, "q_a"), (weight, "weight")): + assert not operands_overlap(out, operand), f"out must not share memory with {name}" + + +q_b_fp8_rope = GlmQBFp8Rope() diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_draft_hidden_half.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_draft_hidden_half.py new file mode 100644 index 000000000000..94966647d7c4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_draft_hidden_half.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm_moe_dsa tp16 draft_hidden_half entry.""" + +import os +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.draft_hidden_half import ( # noqa: E501 + draft_hidden_half as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the glm_moe_dsa tp16 entries are certified on sm_107" +) + + +def _build(spec: dict, seed: int) -> dict: + """The op's arguments for one cell. + + The hidden state has the scale of a residual stream (standard deviation 8); the norm weight + sits around 1; the weight is scaled to unit-variance outputs. `out` starts as NaN, so an + output the kernel never stores fails the comparison. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + rows, k, n = spec["rows"], spec["k"], spec["n"] + x = (8.0 * torch.randn(rows, k, generator=g, device="cuda")).to(torch.bfloat16) + norm_weight = (1.0 + 0.1 * torch.randn(k, generator=g, device="cuda")).to(torch.bfloat16) + weight = (torch.randn(n, k, generator=g, device="cuda") * k**-0.5).to(torch.bfloat16) + in_kernel_norm = spec["in_kernel_norm"] + if not in_kernel_norm: + # The fallback's input: the row(s) already normalized and rounded to BF16. + xf = x.float() + x = (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + spec["eps"])).to(torch.bfloat16) + return dict( + x=x, + weight=weight, + out=torch.full((rows, n), float("nan"), dtype=torch.float32, device="cuda"), + norm_weight=norm_weight if in_kernel_norm else None, + eps=spec["eps"] if in_kernel_norm else None, + ) + + +def _compare(out: torch.Tensor, ref: torch.Tensor, in_kernel_norm: bool) -> None: + """FP32 outputs; the band depends on what can differ. + + Plain GEMM: the same BF16 operands, products exact in FP32, sums in another order (split-K + 8 over a 6144-long K against torch): about 1e-6 of a row's scale. Allowed: 2**-13 of the row + maximum element-wise and 2**-15 relative RMS, two orders above that and orders below a lost + or doubled K slice (about a third of the values). + + In-kernel norm: the kernel normalizes with FlashInfer's summation order and approximate + rsqrt, the reference with torch's, so a normalized element sitting next to a BF16 rounding + midpoint can round the other way (one BF16 ulp; a few FP32 ulps of rstd make that about one + element in 3 * 10**4, so a row sees none or one). One such flip moves every output by at + most 2**-7 of that element's product with its weight, about 2**-10.5 of the row maximum + with these operands. Allowed: 2**-8 of the row maximum element-wise and 2**-10 relative + RMS. A wrong eps, a missing norm weight or an unnormalized row is off by whole values. + """ + if in_kernel_norm: + assert_within_ulp(out, ref, element_ulp=1.0, rms_ulp=0.25, ulp=2.0**-8) + else: + assert_within_ulp(out, ref, element_ulp=1.0, rms_ulp=0.25, ulp=2.0**-13) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + x_before = args["x"].clone() + assert op(**args) is args["out"], "the op writes in place and returns what it wrote" + assert torch.equal(args["x"], x_before) + _compare(args["out"], op.reference(**args), cell.spec["in_kernel_norm"]) + + +def _flashinfer_norm_is_cute_dsl() -> bool: + """Whether trtllm::flashinfer_rmsnorm runs FlashInfer's CuTe DSL kernel, the one the in-kernel + norm reproduces. FlashInfer picks its CUDA kernel under FLASHINFER_USE_CUDA_NORM=1 and where + its own device check says so; without that check the default, CuTe DSL, is assumed.""" + if os.environ.get("FLASHINFER_USE_CUDA_NORM", "0") == "1": + return False + try: + from flashinfer.norm import _use_cuda_norm + except ImportError: + return True + return not _use_cuda_norm(torch.device("cuda", torch.cuda.current_device())) + + +def test_in_kernel_norm_is_flashinfer_rmsnorm_bit_for_bit() -> None: + """The point of the in-kernel norm: the result of FlashInfer's RMSNorm followed by this GEMM, + bit for bit, from one kernel.""" + if not _flashinfer_norm_is_cute_dsl(): + pytest.skip("FlashInfer runs its CUDA RMSNorm kernel here, not the one reproduced") + args = _build(op.CELLS[0].spec, seed=71) + with validating(op): + fused = op(**args) + normed = torch.ops.trtllm.flashinfer_rmsnorm(args["x"], args["norm_weight"], args["eps"]) + plain = op(normed, args["weight"], torch.empty_like(args["out"]), None, None) + assert torch.equal(fused, plain) + + +def test_cuda_graph_capture_and_replay() -> None: + """The draft loop captures the op on the auxiliary stream: it must replay on new contents of + the same buffers. Warm-up compiles the kernel outside the capture.""" + spec = op.CELLS[0].spec + args = _build(spec, seed=72) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + op(**args) + torch.cuda.current_stream().wait_stream(warmup) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(**args) + + args["x"].copy_(_build(spec, seed=73)["x"]) + args["out"].fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + _compare(args["out"], op.reference(**args), in_kernel_norm=True) + + +def test_guard_refuses_eps_without_the_norm_weight() -> None: + """eps alone would be ignored: the raw row would go into the GEMM.""" + args = _build(op.CELLS[0].spec, seed=74) + args["norm_weight"] = None + with pytest.raises(AssertionError, match="pass norm_weight and eps together"): + with validating(op): + op(**args) + + +def test_guard_refuses_an_out_over_the_weight() -> None: + """The owner CTAs store while other clusters still stream the weight: silently wrong.""" + args = _build(op.CELLS[0].spec, seed=75) + rows, n = args["out"].shape + args["out"] = args["weight"].view(torch.float32)[:rows, :n] + with pytest.raises(AssertionError, match="out must not share memory with weight"): + with validating(op): + op(**args) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_indexer_q_tail.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_indexer_q_tail.py new file mode 100644 index 000000000000..c5c26a0a31da --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_indexer_q_tail.py @@ -0,0 +1,147 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm_moe_dsa tp16 indexer_q_tail entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.indexer_q_tail import ( # noqa: E501 + indexer_q_tail as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the glm_moe_dsa tp16 entries are certified on sm_107" +) + +#: GLM-5.2's fused-A output row with the indexer's wk (128) and weights_proj (32) rows: the +#: MLA's 2624 columns, then the indexer K, then the head weights. +_FUSED_A_WIDTH = 2784 +_WEIGHTS_COLUMN = 2624 + 128 +#: Rows of the test's RoPE table, and GLM-5.2's rope_theta. +_POSITIONS = 8192 +_THETA = 8.0e6 +#: An E4M3 NaN: what every Q byte holds until the kernel stores it. +_E4M3_NAN = 0x7F + + +def _indexer_cos_sin(positions: int, rope_dim: int) -> torch.Tensor: + """The indexer's RoPE table: per position the rope_dim / 2 cosines, then the sines.""" + inv_freq = _THETA ** (-torch.arange(0, rope_dim, 2, dtype=torch.float64) / rope_dim) + angle = torch.arange(positions, dtype=torch.float64)[:, None] * inv_freq[None, :] + return torch.cat((angle.cos(), angle.sin()), dim=-1).float().cuda().contiguous() + + +def _build(spec: dict, seed: int) -> dict: + """The op's arguments for one cell. + + q_a and the head weights are views of one fused-A output, as the target passes them. The + weight is scaled so that wq_b's outputs have a standard deviation of 2. The outputs start as + NaN, so anything the kernel does not store fails the comparison. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + m, k = spec["m"], spec["k"] + heads, head_dim = spec["n_heads"], spec["head_dim"] + fused_a = torch.randn(m, _FUSED_A_WIDTH, generator=g, device="cuda").to(torch.bfloat16) + weight = (torch.randn(heads * head_dim, k, generator=g, device="cuda") * 2.0 * k**-0.5).to( + torch.bfloat16 + ) + positions = torch.randint( + 0, _POSITIONS, (m,), generator=g, device="cuda", dtype=torch.int32 + ) + q_fp8 = torch.full((m, heads, head_dim), _E4M3_NAN, dtype=torch.uint8, device="cuda") + return dict( + q_a=fused_a[:, :k], + weight=weight, + positions=positions, + cos_sin=_indexer_cos_sin(_POSITIONS, spec["rope_dim"]), + weights_in=fused_a[:, _WEIGHTS_COLUMN : _WEIGHTS_COLUMN + heads], + q_fp8=q_fp8.view(torch.float8_e4m3fn), + q_scale=torch.full((m, heads, 1), float("nan"), dtype=torch.float32, device="cuda"), + weights_out=torch.full((m, heads), float("nan"), dtype=torch.float32, device="cuda"), + eps=spec["eps"], + n_heads=heads, + head_dim=head_dim, + rope_dim=spec["rope_dim"], + interleave=spec["interleave"], + ue8m0=spec["ue8m0"], + weight_scale=spec["weight_scale"], + ) + + +def _compare(got: tuple, want: tuple) -> None: + """Scales and weights exact where the head's scale agrees; Q within one E4M3 ulp. + + wq_b's FP32 outputs differ from the reference's by the split-K summation order of a + 2048-long dot product and the approximate rsqrt (a few times 1e-6 relative), which changes + the BF16 rounding of about one value in a thousand by one BF16 ulp. Those flips reach the + E4M3 Q only where they cross an E4M3 rounding midpoint (1 in 16 of them) and reach a head's + UE8M0 scale only where the head's amax crosses a power of two (far rarer). Allowed: at most + 2 (token, head) scales differing, exact scaled weights where the scale agrees, every Q value + within one E4M3 ulp of the coarser of the two scales, at most max(8, 0.1 %) Q bytes + differing. A wrong position, pair, sign, scale rule or amax moves whole heads. + """ + q, scale, weights = got + q_ref, scale_ref, weights_ref = want + assert torch.isfinite(scale).all() and torch.isfinite(weights).all(), "an output is unwritten" + q_value = q.float() + assert torch.isfinite(q_value).all(), "a Q byte is unwritten (NaN)" + same_scale = scale == scale_ref + assert int((~same_scale).sum()) <= 2, f"{int((~same_scale).sum())} head scales differ" + same = same_scale.squeeze(-1) + assert torch.equal(weights[same], weights_ref[same]), "the scaled weights differ" + coarser = torch.maximum(scale, scale_ref) + got_value = q_value * scale + want_value = q_ref.float() * scale_ref + magnitude = (torch.maximum(got_value.abs(), want_value.abs()) / coarser).clamp_min(2.0**-6) + ulp = torch.exp2(torch.floor(torch.log2(magnitude)) - 3) * coarser + worst = ((got_value - want_value).abs() / ulp).max().item() + assert worst <= 1.0, f"a Q value is {worst} E4M3 ulp off" + differing = int((q.view(torch.uint8) != q_ref.view(torch.uint8)).sum()) + assert differing <= max(8, q.numel() // 1000), f"{differing} of {q.numel()} Q bytes differ" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + q_a_before, weights_in_before = args["q_a"].clone(), args["weights_in"].clone() + got = op(**args) + assert got[0] is args["q_fp8"] and got[1] is args["q_scale"] and got[2] is args["weights_out"] + assert torch.equal(args["q_a"], q_a_before) + assert torch.equal(args["weights_in"], weights_in_before) + _compare(got, op.reference(**args)) + + +def test_scales_are_powers_of_two() -> None: + """UE8M0: every scale is a power of two at or above amax / 448, so the Q never saturates.""" + args = _build(op.CELLS[0].spec, seed=51) + with validating(op): + q, scale, _ = op(**args) + mantissa = scale.view(torch.int32) & 0x007FFFFF + assert int(mantissa.abs().sum()) == 0, "a scale is not a power of two" + assert float(q.float().abs().max()) <= 448.0 + + +def test_guard_refuses_a_position_past_the_table() -> None: + args = _build(op.CELLS[0].spec, seed=52) + args["positions"][0] = _POSITIONS + with pytest.raises(AssertionError, match="every position must index a row"): + with validating(op): + op(**args) + + +def test_guard_refuses_a_q_over_q_a() -> None: + """The owner CTAs store while other clusters still read q_a: silently wrong.""" + args = _build(op.CELLS[0].spec, seed=53) + m, heads, head_dim = args["q_fp8"].shape + over = args["q_a"].view(torch.uint8)[:, : heads * head_dim].view(m, heads, head_dim) + args["q_fp8"] = over.view(torch.float8_e4m3fn) + with pytest.raises(AssertionError, match="q_fp8 must not share memory with q_a"): + with validating(op): + op(**args) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_low_m_gemm.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_low_m_gemm.py new file mode 100644 index 000000000000..ffe9d1b60579 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_low_m_gemm.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm_moe_dsa tp16 low_m_gemm entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.low_m_gemm import ( # noqa: E501 + low_m_gemm as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the glm_moe_dsa tp16 entries are certified on sm_107" +) + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """One GEMM's worth of accumulation-order difference. + + Kernel and reference multiply the same BF16 operands exactly in FP32 and differ only in the + order of the FP32 sums: the kernel's per-slice tensor-core accumulation and its fixed DSMEM + reduction of the K slices against torch's FP32 matmul. That is about 1e-6 relative to a + row's scale, which moves the one BF16 rounding of an output only where the FP32 sum sits + next to a rounding midpoint: rare one-ulp flips. One BF16 ulp of the row maximum is 2 units + of the band's 2**-8, so 4 units element-wise admits any such flip with margin, and 2 units of + relative RMS sits far above what a handful of flips adds. A lost or doubled K slice, a + wrong tile or a missing tail is off by the size of the values themselves, hundreds of units. + """ + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0) + + +def _build(spec: dict, seed: int): + """Operands of one cell: unit-variance activations, weights scaled to unit-variance outputs. + + `out` starts as NaN, so an output the kernel never stores fails the comparison. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + m, k, n = spec["m"], spec["k"], spec["n"] + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + weight = (torch.randn(n, k, generator=g, device="cuda") * k**-0.5).to(torch.bfloat16) + out = torch.full((m, n), float("nan"), dtype=torch.bfloat16, device="cuda") + return x, weight, out + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + x, weight, out = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + x_before, weight_before = x.clone(), weight.clone() + # Identity, not equality: the op writes the caller's buffer and returns it. + assert op(x, weight, out) is out, "the op writes in place and returns what it wrote" + assert torch.equal(x, x_before) and torch.equal(weight, weight_before) + _compare(out, op.reference(x, weight, out)) + + +def test_cuda_graph_capture_and_replay() -> None: + """The target captures the whole decode step: the op must replay on new contents of the + same buffers. Warm-up compiles the kernel outside the capture.""" + spec = op.CELLS[0].spec + x, weight, out = _build(spec, seed=31) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + op(x, weight, out) + torch.cuda.current_stream().wait_stream(warmup) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(x, weight, out) + + fresh, _, _ = _build(spec, seed=32) + x.copy_(fresh) + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + _compare(out, op.reference(x, weight, out)) + + +def _overlapping(spec: dict, operand: str): + """A cell's operands with `out` placed over `operand`'s memory.""" + x, weight, _ = _build(spec, seed=33) + m, n = spec["m"], spec["n"] + source = x if operand == "x" else weight + buffer = torch.empty(source.numel() + m * n, dtype=torch.bfloat16, device="cuda") + buffer[: source.numel()].copy_(source.view(-1)) + placed = buffer[: source.numel()].view(source.shape) + out = buffer[: m * n].view(m, n) + return (placed, weight, out) if operand == "x" else (x, placed, out) + + +@pytest.mark.parametrize("operand", ["x", "weight"]) +def test_guard_refuses_an_out_over_an_input(operand) -> None: + """The owner CTAs store while other clusters still read the inputs: silently wrong.""" + x, weight, out = _overlapping(op.CELLS[0].spec, operand) + with pytest.raises(AssertionError, match=f"out must not share memory with {operand}"): + with validating(op): + op(x, weight, out) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_o_proj_publish.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_o_proj_publish.py new file mode 100644 index 000000000000..6a856d09fe8e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_o_proj_publish.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm_moe_dsa tp16 o_proj_publish entry. + +One GPU stands in for the group: `lamport` is a plain device buffer laid out like an MNNVL +one-shot workspace (three Lamport buffers, every element the idle -0.0) and `flags` a plain int32 +tensor, so the test can read back exactly which slots the op stored and that it stored nothing +else. On a real group `lamport` is the multicast address and the same stores reach every rank. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.low_m_gemm import ( # noqa: E501 + low_m_gemm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.o_proj_publish import ( # noqa: E501 + o_proj_publish as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the glm_moe_dsa tp16 entries are certified on sm_107" +) + +#: Lamport buffers of a one-shot workspace. +_BUFFERS = 3 +#: Bytes past the slots in each buffer: a store past the end of the slot block lands there. +_SLACK_BYTES = 4096 +#: The BF16 bit pattern of -0.0, the idle value of a Lamport buffer, as int16. +_IDLE = -32768 + + +def _build(spec: dict, seed: int) -> dict: + """The op's arguments for one cell: unit-variance attention outputs, a weight scaled to + unit-variance outputs, a workspace with the slots of every token and rank plus slack.""" + g = torch.Generator(device="cuda").manual_seed(seed) + m, k, n, world = spec["m"], spec["k"], spec["n"], spec["world_size"] + bytes_per_buffer = m * world * n * 2 + _SLACK_BYTES + current = spec["current"] + flags = torch.zeros(16, dtype=torch.int32, device="cuda") + flags[:4] = torch.tensor([current, (current + 2) % _BUFFERS, bytes_per_buffer, 1]) + return dict( + x=torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16), + weight=(torch.randn(n, k, generator=g, device="cuda") * k**-0.5).to(torch.bfloat16), + out=torch.empty(m, n, dtype=torch.bfloat16, device="cuda"), + lamport=torch.full( + (_BUFFERS * bytes_per_buffer // 2,), -0.0, dtype=torch.bfloat16, device="cuda" + ), + flags=flags, + world_size=world, + rank=spec["rank"], + ) + + +def _slots(args: dict) -> tuple[torch.Tensor, torch.Tensor]: + """(this rank's [M, N] slots in the current buffer, a mask of every element of `lamport` + outside them).""" + m, n = args["x"].shape[0], args["weight"].shape[0] + world, rank = args["world_size"], args["rank"] + current, _, bytes_per_buffer, _ = args["flags"][:4].tolist() + buffer = args["lamport"].view(_BUFFERS, bytes_per_buffer // 2)[current] + slots = buffer[: m * world * n].view(m, world, n) + elsewhere = torch.ones(args["lamport"].numel(), dtype=torch.bool, device="cuda") + mask = elsewhere.view(_BUFFERS, -1)[current][: m * world * n].view(m, world, n) + mask[:, rank, :] = False + return slots[:, rank, :], elsewhere + + +def _compare(published: torch.Tensor, ref: torch.Tensor) -> None: + """One GEMM's worth of accumulation-order difference, as for low_m_gemm. + + The FP32 sums of the kernel (split-K 2 over a 1024-long K) and of torch differ in order by + about 1e-6 relative, which flips the BF16 rounding of a rare output by one ulp: 2 units of + the band's 2**-8 at the row maximum, so 4 units element-wise and 2 units of relative RMS. A + value in the wrong slot, a lost K slice or a wrong rounding is off by whole values. + """ + assert_within_ulp(published, ref, element_ulp=4.0, rms_ulp=2.0) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + assert op(**args) is None + published, elsewhere = _slots(args) + assert bool((published.view(torch.int16) != _IDLE).all()), "a slot holds -0.0 (unwritten)" + idle = args["lamport"].view(torch.int16)[elsewhere] + assert bool((idle == _IDLE).all()), "the op stored outside this rank's slots" + _compare(published, op.reference(**args)) + + +def test_published_values_are_the_gemm_output_bit_for_bit() -> None: + """The publish epilogue rounds as the plain epilogue does: the slots hold low_m_gemm's + output at the same tactic, -0.0 written as +0.0.""" + args = _build(op.CELLS[0].spec, seed=61) + with validating(op): + op(**args) + published, _ = _slots(args) + plain = low_m_gemm(args["x"], args["weight"], torch.empty_like(args["out"])) + plain = torch.where(plain == 0, torch.zeros_like(plain), plain) + assert torch.equal(published.view(torch.int16), plain.view(torch.int16)) + + +def test_cuda_graph_capture_and_replay() -> None: + """Captured once, the op publishes into whichever buffer the flags mark at replay.""" + args = _build(op.CELLS[0].spec, seed=62) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + op(**args) + torch.cuda.current_stream().wait_stream(warmup) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(**args) + + args["lamport"].fill_(-0.0) + args["flags"][0] = (int(args["flags"][0]) + 1) % _BUFFERS + args["x"].copy_(_build(op.CELLS[0].spec, seed=63)["x"]) + graph.replay() + torch.cuda.synchronize() + published, elsewhere = _slots(args) + assert bool((args["lamport"].view(torch.int16)[elsewhere] == _IDLE).all()) + _compare(published, op.reference(**args)) + + +def test_guard_refuses_a_rank_outside_the_group() -> None: + """Rank 16 of 16 lands in the next token's slots.""" + args = _build(op.CELLS[0].spec, seed=64) + args["rank"] = args["world_size"] + with pytest.raises(AssertionError, match="outside a group of"): + with validating(op): + op(**args) + + +def test_guard_refuses_a_buffer_smaller_than_the_slots() -> None: + """The slots past the buffer's end land in the next Lamport buffer.""" + args = _build(op.CELLS[0].spec, seed=65) + m, n = args["x"].shape[0], args["weight"].shape[0] + args["flags"][2] = m * args["world_size"] * n * 2 - 16 + with pytest.raises(AssertionError, match="smaller than the"): + with validating(op): + op(**args) + + +def test_guard_refuses_a_buffer_index_out_of_range() -> None: + args = _build(op.CELLS[0].spec, seed=66) + args["flags"][0] = _BUFFERS + with pytest.raises(AssertionError, match="current buffer index"): + with validating(op): + op(**args) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_q_b_fp8_rope.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_q_b_fp8_rope.py new file mode 100644 index 000000000000..cd70e5d27797 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_q_b_fp8_rope.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm_moe_dsa tp16 q_b_fp8_rope entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.q_b_fp8_rope import ( # noqa: E501 + q_b_fp8_rope as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the glm_moe_dsa tp16 entries are certified on sm_107" +) + +#: GLM-5.2's fused-A output row: q_lora_rank 2048, kv_lora_rank 512, qk_rope_head_dim 64. +_FUSED_A_WIDTH = 2624 +#: Rows of the test's RoPE table, and GLM-5.2's rope_theta. +_POSITIONS = 8192 +_THETA = 8.0e6 +#: An E4M3 NaN: what every output byte holds until the kernel stores it. +_E4M3_NAN = 0x7F + + +def _mla_cos_sin(positions: int, rope_dim: int) -> torch.Tensor: + """The MLA RoPE table in the layout the attention backend builds for MLA: per position + rope_dim (cos, sin) pairs, the rope_dim / 2 frequencies theta ** (-2i / rope_dim) twice.""" + inv_freq = _THETA ** (-torch.arange(0, rope_dim, 2, dtype=torch.float64) / rope_dim) + angle = torch.arange(positions, dtype=torch.float64)[:, None] * inv_freq[None, :] + angle = torch.cat((angle, angle), dim=-1) + return torch.stack((angle.cos(), angle.sin()), dim=-1).float().cuda().contiguous() + + +def _build(spec: dict, seed: int) -> dict: + """The op's arguments for one cell. + + q_a is the q_lora_rank-column row view of a fused-A output, as the target passes it. The + weight is scaled so that the FP32 values before the conversion have a standard deviation of + 4 (times q_scale): E4M3 normals of every exponent from 2**-6 up, and subnormals below. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + m, k, seq_len = spec["m"], spec["k"], spec["seq_len"] + n = spec["heads"] * spec["head_dim"] + fused_a = torch.randn(m, _FUSED_A_WIDTH, generator=g, device="cuda").to(torch.bfloat16) + weight = (torch.randn(n, k, generator=g, device="cuda") * 4.0 * k**-0.5).to(torch.bfloat16) + kv_lens = torch.randint( + seq_len, + _POSITIONS + 1, + (m // seq_len,), + generator=g, + device="cuda", + dtype=torch.int32, + ) + return dict( + q_a=fused_a[:, :k], + weight=weight, + out=torch.full((m, n), _E4M3_NAN, dtype=torch.uint8, device="cuda"), + kv_lens=kv_lens, + cos_sin=_mla_cos_sin(_POSITIONS, spec["rope_dim"]), + q_scale=torch.full((1,), spec.get("q_scale", 1.0), dtype=torch.float32, device="cuda"), + eps=spec["eps"], + head_dim=spec["head_dim"], + rope_dim=spec["rope_dim"], + seq_len=seq_len, + plan=spec["plan"], + ) + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Every E4M3 value within one E4M3 ulp, and only the rare ones not equal. + + Before its single rounding, the kernel's FP32 value differs from the reference's by the + split-K summation order of a 2048-long dot product (about 3e-6 relative), the approximate + rsqrt of the row RMS (about 2**-22) and the contraction of the RoPE products into FMAs (one + FP32 ulp): a few times 1e-6 relative. E4M3 spaces its values 2**-4 to 2**-3 apart + (relative), so a value converts differently only within that distance of a rounding + midpoint: a probability of about 1e-4 per element, about two of the 18432 values of a verify + cell. Allowed: one E4M3 ulp per element, and at most max(4, 0.1 %) differing bytes. A wrong + position, pair, sign or scale moves most RoPE values by many ulps, and a lost K slice moves + every value. + """ + got = out.view(torch.float8_e4m3fn).float() + want = ref.view(torch.float8_e4m3fn).float() + assert torch.isfinite(got).all(), "the op left an output byte unwritten (NaN) or wrote NaN" + magnitude = torch.maximum(got.abs(), want.abs()).clamp_min(2.0**-6) + ulp = torch.exp2(torch.floor(torch.log2(magnitude)) - 3) + worst = ((got - want).abs() / ulp).max().item() + assert worst <= 1.0, f"an E4M3 value is {worst} ulp off" + differing = int((out != ref).sum()) + assert differing <= max(4, out.numel() // 1000), f"{differing} of {out.numel()} bytes differ" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + q_a_before, weight_before = args["q_a"].clone(), args["weight"].clone() + assert op(**args) is args["out"], "the op writes in place and returns what it wrote" + assert torch.equal(args["q_a"], q_a_before) and torch.equal(args["weight"], weight_before) + _compare(args["out"], op.reference(**args)) + + +def test_q_scale_is_applied_and_saturates() -> None: + """A scale of 32 pushes the values past 448: the conversion saturates instead of producing + NaN, and the scale multiplies before it (c1 runs at 1, where the scale is invisible).""" + spec = dict(op.CELLS[0].spec, q_scale=32.0) + args = _build(spec, seed=41) + with validating(op): + op(**args) + got = args["out"].view(torch.float8_e4m3fn).float() + assert (got.abs() == 448.0).any(), "no value reached the saturation bound" + _compare(args["out"], op.reference(**args)) + + +def test_plans_agree() -> None: + """The plans change the launch and the FP32 summation order, not the definition.""" + spec = op.CELLS[0].spec + results = {} + for plan in ("regular", "beside_topk", "early_gather"): + args = _build(dict(spec, plan=plan), seed=43) + op(**args) + results[plan] = args["out"] + _compare(results["beside_topk"], results["regular"]) + _compare(results["early_gather"], results["regular"]) + + +def test_cuda_graph_capture_and_replay() -> None: + """The target captures the whole decode step: the op must replay on new contents of the same + buffers. Warm-up compiles the kernel and builds its placeholders outside the capture.""" + spec = op.CELLS[0].spec + args = _build(spec, seed=44) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + op(**args) + torch.cuda.current_stream().wait_stream(warmup) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(**args) + + fresh = _build(spec, seed=45) + args["q_a"].copy_(fresh["q_a"]) + args["kv_lens"].copy_(fresh["kv_lens"]) + args["out"].fill_(_E4M3_NAN) + graph.replay() + torch.cuda.synchronize() + _compare(args["out"], op.reference(**args)) + + +def test_guard_refuses_a_table_of_another_row_width() -> None: + """A table with rope_dim / 2 pairs per position (no repeated half) is read at the wrong rows.""" + args = _build(op.CELLS[0].spec, seed=46) + args["cos_sin"] = args["cos_sin"][:, : args["rope_dim"] // 2].contiguous() + with pytest.raises(AssertionError, match="cos_sin must be"): + with validating(op): + op(**args) + + +def test_guard_refuses_a_position_past_the_table() -> None: + args = _build(op.CELLS[0].spec, seed=47) + args["kv_lens"].fill_(_POSITIONS + 1) + with pytest.raises(AssertionError, match="cos/sin table must hold every position"): + with validating(op): + op(**args) + + +def test_guard_refuses_kv_lens_shorter_than_the_step() -> None: + """kv_lens below seq_len gives negative positions: reads before the table.""" + args = _build(op.CELLS[0].spec, seed=48) + args["kv_lens"].fill_(args["seq_len"] - 1) + with pytest.raises(AssertionError, match="kv_lens must cover the step's own tokens"): + with validating(op): + op(**args) + + +def test_guard_refuses_an_out_over_q_a() -> None: + """The owner CTAs store while other clusters still read q_a: silently wrong.""" + args = _build(op.CELLS[0].spec, seed=49) + m, n = args["out"].shape + args["out"] = args["q_a"].view(torch.uint8)[:, :n] + assert args["out"].shape == (m, n) + with pytest.raises(AssertionError, match="out must not share memory with q_a"): + with validating(op): + op(**args) From 44d3f2725f5c90740dd2c145ac6c0aac0d1223d9 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:04:23 -0700 Subject: [PATCH 14/53] Let the TRTLLM-Gen NVFP4 MoE op wait for its input right before FC1 fp4_block_scale_moe_runner takes an optional hidden_states_ready_event: the raw cudaEvent_t (torch.cuda.Event.cuda_event) of an event recorded on the stream that writes hidden_states and hidden_states_scale. MoE::Runner::run waits on it with cudaStreamWaitEvent right before the FC1 launch, the first kernel that reads them, so the routing kernel launched before it depends on the router alone while the input quantization runs on another stream. Without the argument nothing waits and the op launches what it launched before. Signed-off-by: Dongfeng Yu --- .../kernels/moe/trtllmGen/runner.cu | 5 +++++ cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.h | 3 +++ .../thop/moe/trtllmGen/fp4BlockScaleMoe.cpp | 12 +++++++++--- .../_torch/custom_ops/trtllm_gen_custom_ops.py | 17 +++++++++++++---- 4 files changed, 30 insertions(+), 7 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.cu b/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.cu index 01ec7600681c..ab5c9e5d945d 100644 --- a/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.cu +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.cu @@ -840,6 +840,11 @@ void Runner::run( auto const& config = mPassingConfigs[configIndex]; + if (args.hidden_states_ready != nullptr) + { + // gemm1 is the first kernel that reads hidden_states; routing, launched before this point, does not wait. + TLLM_CUDA_CHECK(cudaStreamWaitEvent(stream, args.hidden_states_ready, 0)); + } mPermuteGemm1.run(args.hidden_states, hidden_states_scale_linear, args.gemm1_weights, args.gemm1_weights_scale, workspace.expert_weights, args.output1_scales_scalar, args.output1_scales_gate_scalar, args.gemm1_bias, args.gemm1_alpha, args.gemm1_beta, args.gemm1_clamp_limit, workspace.gemm1_output, workspace.gemm1_output_scale, diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.h b/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.h index 453a29671f4a..8f7ee3a84979 100644 --- a/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.h +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.h @@ -288,6 +288,9 @@ struct MoERunnerArgs // [hidden_size/128, num_tokens] in float for e4m3 DS recipe // and [num_tokens, hidden_size/16] in float for e2m1 void* hidden_states_scale = nullptr; + // Optional: recorded on another stream once hidden_states and hidden_states_scale are written. The runner waits + // on it right before gemm1, their first reader, so the routing launched before gemm1 does not wait for them. + cudaEvent_t hidden_states_ready = nullptr; // Optional inputs: void* topk_weights = nullptr; // [num_tokens, top_k] with quantized weights diff --git a/cpp/tensorrt_llm/thop/moe/trtllmGen/fp4BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/moe/trtllmGen/fp4BlockScaleMoe.cpp index fac7cfb36e77..d054b29b273a 100644 --- a/cpp/tensorrt_llm/thop/moe/trtllmGen/fp4BlockScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/moe/trtllmGen/fp4BlockScaleMoe.cpp @@ -46,7 +46,8 @@ std::vector run_fp4_block_scale_moe_runner(torch::optional const routed_scaling_factor, int64_t const tile_tokens_dim, int64_t const routing_method_type, bool const do_finalize, btg::Dtype const dtype, MoeRunnerType& moe_runner, int64_t const moeConfigIndex, torch::optional const& topk_weights, torch::optional const& topk_ids, - torch::optional const& out_tensor = torch::nullopt) + torch::optional const& out_tensor = torch::nullopt, + std::optional const hidden_states_ready_event = std::nullopt) { TORCH_CHECK(dtype == btg::Dtype::E4m3 || dtype == btg::Dtype::E2m1, "dtype can only be e4m3 or e2m1."); TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "Only SM100f is supported by FP4 block scale MOE"); @@ -154,6 +155,10 @@ std::vector run_fp4_block_scale_moe_runner(torch::optional(static_cast(hidden_states_ready_event.value())) + : nullptr; args.topk_weights = topk_weights.has_value() ? topk_weights.value().data_ptr() : nullptr; args.topk_ids = topk_ids.has_value() ? static_cast(topk_ids.value().data_ptr()) : nullptr; @@ -589,7 +594,8 @@ class FP4BlockScaleMoeRunner : public torch::CustomClassHolder int64_t const local_expert_offset, int64_t const local_num_experts, std::optional const routed_scaling_factor, int64_t const routing_method_type, bool const do_finalize, std::vector moeConfigIndex, torch::optional const& topk_weights, - torch::optional const& topk_ids, torch::optional const& output = torch::nullopt) + torch::optional const& topk_ids, torch::optional const& output = torch::nullopt, + std::optional const hidden_states_ready_event = std::nullopt) { int64_t tileN = moeConfigIndex[0]; int64_t config = moeConfigIndex[1]; @@ -639,7 +645,7 @@ class FP4BlockScaleMoeRunner : public torch::CustomClassHolder gemm2_weights_scale, gemm2_bias, output1_scales_scalar, output1_scales_gate_scalar, output2_scales_scalar, num_experts, top_k, n_group, topk_group, intermediate_size, local_expert_offset, local_num_experts, routed_scaling_factor, tileN, routing_method_type, do_finalize, mDtypeElt, *runner, configToRun, - topk_weights, topk_ids, output); + topk_weights, topk_ids, output, hidden_states_ready_event); } private: diff --git a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py index 0063abfa4f1c..053d55fc50f1 100644 --- a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py @@ -499,6 +499,7 @@ def forward( inputs: List[torch.Tensor], tactic: List[int] = [-1, -1], output: Optional[torch.Tensor] = None, + hidden_states_ready_event: Optional[int] = None, ) -> torch.Tensor: assert isinstance(tactic, list) @@ -516,7 +517,8 @@ def forward( self.n_group, self.topk_group, self.intermediate_size, self.local_expert_offset, self.local_num_experts, self.routed_scaling_factor, self.routing_method_type, - self.do_finalize, tactic, args.topk_weights, args.topk_ids, output) + self.do_finalize, tactic, args.topk_weights, args.topk_ids, output, + hidden_states_ready_event) def get_valid_tactics(self, inputs: List[torch.Tensor], profile: OptimizationProfile, @@ -679,7 +681,12 @@ def fp4_block_scale_moe_runner(routing_logits: Optional[torch.Tensor], topk_ids: Optional[torch.Tensor] = None, output: Optional[torch.Tensor] = None, tune_max_num_tokens: int = 8192, - use_dp: bool = False) -> List[torch.Tensor]: + use_dp: bool = False, + hidden_states_ready_event: Optional[int] = None + ) -> List[torch.Tensor]: + """hidden_states_ready_event: raw cudaEvent_t (torch.cuda.Event.cuda_event) + recorded on another stream once hidden_states and hidden_states_scale are + written; the op waits on it right before FC1, after launching routing.""" tuner = AutoTuner.get() kernel_runner = FP4BlockScaleMoERunner( @@ -755,7 +762,8 @@ def fp4_block_scale_moe_runner(routing_logits: Optional[torch.Tensor], result = kernel_runner( input_tensors, tactic=[-1, -1] if best_tactic == -1 else best_tactic, - output=output) + output=output, + hidden_states_ready_event=hidden_states_ready_event) # When output is provided and do_finalize=True, the result is written in-place to output. # Return empty tensor to avoid aliasing constraint violation in PyTorch 2.9.1+ # (custom op output cannot be the same tensor as input). @@ -823,7 +831,8 @@ def _(routing_logits, topk_ids: Optional[torch.Tensor] = None, output: Optional[torch.Tensor] = None, tune_max_num_tokens: int = 8192, - use_dp: bool = False) -> List[torch.Tensor]: + use_dp: bool = False, + hidden_states_ready_event: Optional[int] = None) -> List[torch.Tensor]: if do_finalize: num_tokens = hidden_states.shape[0] hidden_size = hidden_states.shape[1] * 2 From 712292cdae69fb9fb5eaa222431f1f69a0545ff8 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:04:34 -0700 Subject: [PATCH 15/53] Compute the router GEMV inside the TRTLLM-Gen MoE routing kernel on request fp4_block_scale_moe_runner takes optional router_hidden_states (bf16 [T, H]), router_weight (bf16 [E, H]) and router_workspace. With them, DeepSeek-V3 routing without groups runs routingFusedRouterKernel instead of the routing kernel: 128 CTAs compute the logits of dsv3_router_gemm_op bit for bit (the same per-thread FMA chains and butterfly order) and publish slice sums to a workspace exchange buffer whose idle value is -0.0; one CTA polls them, runs the small-token kernel's routing (sigmoid + bias, top-8, scaled sum normalization, same operations in the same order) and writes every output that kernel writes. routing_logits then only gives the shape. The kernel covers 256 experts, top-8, hidden 6144 and 1-8 tokens and refuses other shapes. make_fused_router_workspace builds the workspace; two exchange buffers alternate by an epoch word, so launches that share a workspace must run in stream order. Without the router arguments the op routes as before. PDL: before its wait the kernel reads only the router weight and the routing bias. The routing CTA triggers its dependents after all routing outputs are written; the GEMV CTAs exit without triggering. Signed-off-by: Dongfeng Yu --- .../trtllmGen/routing/RoutingFusedRouter.cu | 520 ++++++++++++++++++ .../moe/trtllmGen/routing/RoutingKernel.h | 15 + .../kernels/moe/trtllmGen/runner.cu | 38 ++ .../kernels/moe/trtllmGen/runner.h | 11 + .../thop/moe/trtllmGen/fp4BlockScaleMoe.cpp | 67 ++- .../custom_ops/trtllm_gen_custom_ops.py | 51 +- 6 files changed, 684 insertions(+), 18 deletions(-) create mode 100644 cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingFusedRouter.cu diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingFusedRouter.cu b/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingFusedRouter.cu new file mode 100644 index 000000000000..329225b1021f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingFusedRouter.cu @@ -0,0 +1,520 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Router GEMV + DeepSeek-style routing (sigmoid + bias, top-8 of 256 experts, scaled sum normalization) in one kernel +// for 1-8 tokens, with the outputs of routingIndicesSmallTokenKernel. +// +// Logits. Every logit is the value dsv3RouterGemm.cu's router_gemm_kernel computes (bit for bit). There, thread t +// (0..127) of the CTA of expert e accumulates acc = fma(a, w, acc) over k = ki * 1024 + 8 t + j (ki = 0..5, j = 0..7), +// in that order; warp s (the "slice" of threads 32 s .. 32 s + 31) adds its lanes with the xor butterfly 16, 8, 4, 2, 1 +// (lane 0's value), and the four slice sums are added to 0.0f in slice order. Here a GEMV CTA covers one slice of 8 +// experts: it stages the slice's activations of every token in shared memory (24 KB at 8 tokens), runs the same +// per-thread chains (fma.rn.f32.bf16 computes fma(float(a), float(w), acc) without conversion instructions), adds the +// 32 positions with the same pairs as the butterfly and publishes the slice sums. 128 CTAs cover 256 experts x 4 +// slices, and each CTA stages a quarter of the activations instead of all of them. +// +// Handoff. The slice sums go to one of two exchange buffers in the workspace (by the parity of the epoch word) whose +// idle value is -0.0. A sum of -0.0 is published as +0.0, which gives the same logit, since the four sums are added to +// 0.0f. The routing CTA polls its token rows until no -0.0 is left, so no fence or counter is needed. It resets the +// other buffer, consumed by the previous launch, while it waits, and advances the epoch at its end. Launches that share +// a workspace must therefore run in stream order. +// +// Routing. Phase A of routingIndicesSmallTokenKernel (one warp per token: sigmoid + bias, the packed-key top-K, the +// scaled sum normalization; the same operations in the same order). The metadata then follows from the set of +// selected local experts after a single barrier: with no more tokens than tile rows every selected local expert owns +// exactly one CTA tile, so its tile index is its rank among the selected local experts and its rows start at +// rank * tile, which is what the small-token kernel's scans compute. + +#include "RoutingCustomPolicy.cuh" +#include "RoutingKernel.h" + +#include + +namespace moe::dev::routing +{ +namespace routingCustom +{ +namespace +{ + +using FusedExpertSelect = TopKExpertSelect; +using FusedKernelParams = KernelParams; + +constexpr int kNumExperts = 256; +constexpr int kTopK = 8; +constexpr int kMaxTokens = 8; +constexpr int kHiddenDim = 6144; +constexpr int kEltsPerVec = 8; // bf16 per 16-byte load, as in the router GEMV +constexpr int kKIterElts = kEltsPerVec * 128; // elements per k-iteration of the router GEMV +constexpr int kKIters = kHiddenDim / kKIterElts; // 6 +constexpr int kSlices = 4; // warps per expert in the router GEMV +constexpr int kExpertsPerCta = 8; +constexpr int kNumGemvCtas = kNumExperts / kExpertsPerCta * kSlices; // 128 +constexpr int kNumThreads = 256; +constexpr int kSliceVecs = kKIters * WarpSize; // 16-byte vectors of one token's slice +constexpr int kActBytes = kMaxTokens * kSliceVecs * 16; // 24 KB +constexpr int kPartBytes = WarpSize * kExpertsPerCta * kMaxTokens * 4; // 8 KB +constexpr int kSmemBytes = kActBytes + kPartBytes; +constexpr int kExchangeFloats = kMaxTokens * kNumExperts * kSlices; +constexpr int kKIdxRowBytes = 8; +constexpr uint32_t kNegZeroBits = 0x80000000u; + +static_assert(kNumExperts == kNumThreads, "the routing CTA runs one thread per expert"); +static_assert(2 * kExchangeFloats + 4 == FusedRouterWorkspaceWords, "workspace layout"); + +struct FusedRouterParams +{ + FusedKernelParams routing; + __nv_bfloat16 const* hiddenStates; + __nv_bfloat16 const* routerWeight; + int32_t* workspace; +}; + +__device__ __forceinline__ uint32_t smemAddr(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ void gridDependencyWait() +{ + // The memory clobber keeps the compiler from moving memory accesses across the wait. + asm volatile("griddepcontrol.wait;" ::: "memory"); +} + +// Non-coherent 16-byte load of router weights. +__device__ __forceinline__ uint4 loadWeight16(void const* ptr) +{ + uint4 v; + asm volatile("ld.global.nc.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(ptr)); + return v; +} + +__device__ __forceinline__ uint32_t loadVolatile(uint32_t const* ptr) +{ + uint32_t v; + asm volatile("ld.volatile.global.u32 %0, [%1];" : "=r"(v) : "l"(ptr) : "memory"); + return v; +} + +__device__ __forceinline__ float4 loadVolatile4(float const* ptr) +{ + float4 v; + asm volatile("ld.volatile.global.v4.f32 {%0, %1, %2, %3}, [%4];" + : "=f"(v.x), "=f"(v.y), "=f"(v.z), "=f"(v.w) + : "l"(ptr) + : "memory"); + return v; +} + +// A slice sum for the routing CTA, which polls it with volatile loads. +__device__ __forceinline__ void storeRelaxed(float* ptr, float v) +{ + asm volatile("st.relaxed.gpu.global.f32 [%0], %1;" ::"l"(ptr), "f"(v) : "memory"); +} + +__device__ __forceinline__ bool isNegZero(float v) +{ + return __float_as_uint(v) == kNegZeroBits; +} + +// acc = fma(a[j], w[j], acc) for j = 0..7, in that order, on 8 packed bf16 pairs. +__device__ __forceinline__ void fmaBf16x8(float& acc, uint4 const& a, uint4 const& w) +{ + uint32_t const aw[4] = {a.x, a.y, a.z, a.w}; + uint32_t const ww[4] = {w.x, w.y, w.z, w.w}; +#pragma unroll + for (int q = 0; q < 4; q++) + { + asm("{\n" + ".reg .b16 al, ah, bl, bh;\n" + "mov.b32 {al, ah}, %1;\n" + "mov.b32 {bl, bh}, %2;\n" + "fma.rn.f32.bf16 %0, al, bl, %0;\n" + "fma.rn.f32.bf16 %0, ah, bh, %0;\n" + "}\n" + : "+f"(acc) + : "r"(aw[q]), "r"(ww[q])); + } +} + +// The slice GEMV of one CTA: slice r of experts 8 g .. 8 g + 7. Thread mapping: warp w covers positions 8 (w & 3) .. +// + 7 of the slice and experts 4 (w >> 2) .. + 3, so the 4 expert groups of a warp read the same 128 bytes of +// activations. Returns through emit(eLocal, m, sliceSum) once per (expert, token). +template +__device__ __forceinline__ void sliceGemv(uint8_t* smem, __nv_bfloat16 const* hiddenStates, uint4 const (&w)[kKIters], + int r, int t, int eLocal, Emit const& emit) +{ + uint4* act = reinterpret_cast(smem); + float* part = reinterpret_cast(smem + kActBytes); + + // Stage the slice: token m, k-iteration ki, positions 32 r .. 32 r + 31. + constexpr int kVecs = kNumTokens * kSliceVecs; +#pragma unroll + for (int j = 0; j < (kVecs + kNumThreads - 1) / kNumThreads; ++j) + { + int const c = threadIdx.x + kNumThreads * j; + if (c < kVecs) + { + int const m = c / kSliceVecs; + int const rem = c - m * kSliceVecs; + __nv_bfloat16 const* src = hiddenStates + m * kHiddenDim + (rem / WarpSize) * kKIterElts + + (WarpSize * r + rem % WarpSize) * kEltsPerVec; + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(smemAddr(&act[c])), "l"(src) : "memory"); + } + } + asm volatile("cp.async.commit_group;\ncp.async.wait_group 0;" ::: "memory"); + __syncthreads(); + + // The per-thread chains of router thread 32 r + t for one expert and every token. + float acc[kMaxTokens] = {}; +#pragma unroll + for (int ki = 0; ki < kKIters; ki++) + { +#pragma unroll + for (int m = 0; m < kNumTokens; m++) + { + fmaBf16x8(acc[m], act[(m * kKIters + ki) * WarpSize + t], w[ki]); + } + } + + // The butterfly of the slice: 4 threads per (expert, token) fold the positions {sub, sub + 4, ..., sub + 28} + // through levels 16, 8 and 4, two shuffles add levels 2 and 1. + float4* row = reinterpret_cast(part + (t * kExpertsPerCta + eLocal) * kMaxTokens); + row[0] = make_float4(acc[0], acc[1], acc[2], acc[3]); + row[1] = make_float4(acc[4], acc[5], acc[6], acc[7]); + __syncthreads(); + int const pair = threadIdx.x / 4; + int const sub = threadIdx.x % 4; + int const eL = pair / kMaxTokens; + int const m = pair % kMaxTokens; + auto v = [&](int tt) { return part[(tt * kExpertsPerCta + eL) * kMaxTokens + m]; }; + float const s1a = v(sub) + v(sub + 16); + float const s1b = v(sub + 4) + v(sub + 20); + float const s1c = v(sub + 8) + v(sub + 24); + float const s1d = v(sub + 12) + v(sub + 28); + float sum = (s1a + s1c) + (s1b + s1d); + sum += __shfl_xor_sync(0xffffffff, sum, 2); + sum += __shfl_xor_sync(0xffffffff, sum, 1); + if (sub == 0 && m < kNumTokens) + { + emit(eL, m, sum); + } +} + +template +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) +__global__ void __launch_bounds__(kNumThreads, 1) routingFusedRouterKernel(FusedRouterParams const __grid_constant__ p) +{ + extern __shared__ __align__(128) uint8_t smemSlice[]; + __shared__ __align__(8) uint8_t smemKIdx[kNumExperts * kKIdxRowBytes]; + __shared__ int16_t smemTopKExpert[kMaxTokens * kTopK]; + __shared__ float smemBias[kNumExperts]; + __shared__ __align__(16) uint32_t smemSelected[kNumExperts / 32]; + + auto const& params = p.routing; + float* const buffers = reinterpret_cast(p.workspace); + uint32_t* const epochPtr = reinterpret_cast(p.workspace + 2 * kExchangeFloats); + + if (blockIdx.x < kNumGemvCtas) + { + int const r = blockIdx.x % kSlices; + int const g = blockIdx.x / kSlices; + int const warp = threadIdx.x / WarpSize; + int const lane = threadIdx.x % WarpSize; + int const t = 8 * (warp % 4) + lane % 8; + int const eLocal = 4 * (warp / 4) + lane / 8; + // The weights are not written by the primary grid: load them before the wait. A shared-memory store of their + // XOR ahead of a barrier makes the loads complete before the wait instead of being scheduled after it. + uint4 w[kKIters]; + __nv_bfloat16 const* wRow + = p.routerWeight + (kExpertsPerCta * g + eLocal) * kHiddenDim + (WarpSize * r + t) * kEltsPerVec; + uint32_t loaded = 0; +#pragma unroll + for (int ki = 0; ki < kKIters; ki++) + { + w[ki] = loadWeight16(wRow + ki * kKIterElts); + loaded ^= w[ki].x ^ w[ki].y ^ w[ki].z ^ w[ki].w; + } + reinterpret_cast(smemSlice + kActBytes)[threadIdx.x] = loaded; + __syncthreads(); + gridDependencyWait(); + // The previous launch on this workspace precedes the primary grid in the stream, so it has completed. + float* const exchange = buffers + (loadVolatile(epochPtr) & 1u) * kExchangeFloats; + sliceGemv(smemSlice, p.hiddenStates, w, r, t, eLocal, + [&](int eL, int m, float sliceSum) + { + storeRelaxed(&exchange[(m * kNumExperts + kExpertsPerCta * g + eL) * kSlices + r], + isNegZero(sliceSum) ? 0.f : sliceSum); + }); + // Exit without a trigger: the dependents launch once the routing CTA triggers. + return; + } + + // Routing CTA. Independent of the primary grid: clear this expert's row of the byte map, stage its bias. + *reinterpret_cast(&smemKIdx[threadIdx.x * kKIdxRowBytes]) = uint2{0xFFFFFFFFu, 0xFFFFFFFFu}; + auto const& pre = params.mExpertSelectParams.mPreprocessParams; + smemBias[threadIdx.x] = loadScalar(pre.ptrRoutingBias, threadIdx.x, pre.dtypeBias); + if (threadIdx.x < kNumExperts / 32) + { + smemSelected[threadIdx.x] = 0u; + } + __syncthreads(); + gridDependencyWait(); + uint32_t const epoch = loadVolatile(epochPtr); + float* const exchange = buffers + (epoch & 1u) * kExchangeFloats; + { + // The buffer of the previous launch, which its routing CTA has consumed: idle again for the next launch. + float4* other = reinterpret_cast(buffers + ((epoch + 1u) & 1u) * kExchangeFloats); + for (int i = threadIdx.x; i < kExchangeFloats / 4; i += kNumThreads) + { + other[i] = make_float4(-0.f, -0.f, -0.f, -0.f); + } + } + + // ----- Phase A: one warp per token, as in routingIndicesSmallTokenKernel ----- + int32_t const warpIdx = __shfl_sync(0xffffffff, threadIdx.x / WarpSize, 0); + int32_t const laneIdx = cutlass::arch::LaneId(); + auto block = cg::this_thread_block(); + auto warp = cg::tiled_partition(block); + if (warpIdx < kNumTokens) + { + constexpr int VecSize = kNumExperts / WarpSize; + float4 slices[VecSize]; + bool ready; + do + { + ready = true; +#pragma unroll + for (int i = 0; i < VecSize; ++i) + { + slices[i] = loadVolatile4(exchange + (warpIdx * kNumExperts + i * WarpSize + laneIdx) * kSlices); + ready = ready && !isNegZero(slices[i].x) && !isNegZero(slices[i].y) && !isNegZero(slices[i].z) + && !isNegZero(slices[i].w); + } + } while (!ready); + + float score[VecSize]; + int32_t idx[VecSize]; +#pragma unroll + for (int i = 0; i < VecSize; ++i) + { + idx[i] = i * WarpSize + laneIdx; + float logit = 0.0f; + logit += slices[i].x; + logit += slices[i].y; + logit += slices[i].z; + logit += slices[i].w; + score[i] = sigmoid_accurate(logit) + smemBias[idx[i]]; + } + float topScore[kTopK]; + int32_t topExpert[kTopK]; + topk::reduceTopK(warp, topScore, topExpert, score, idx, float{-INFINITY}, kTopK); + + // Lane k's top-K entry, the value the small-token kernel reads as warpTopKScore[laneIdx]. + float myScore = 0.f; + int32_t myExpert = 0; +#pragma unroll + for (int k = 0; k < kTopK; ++k) + { + myScore = laneIdx == k ? topScore[k] : myScore; + myExpert = laneIdx == k ? topExpert[k] : myExpert; + } + bool const isTop = laneIdx < kTopK; + float const sigmoidScore = isTop ? (myScore - smemBias[myExpert]) : 0.f; + // cg::reduce's butterfly (xor 16, 8, 4, 2, 1) adds +0.0 to lanes 0-7 in its first two levels, which leaves + // them unchanged (no sigmoidScore is -0.0); levels 4, 2 and 1 then pair lanes 0-7 as here. + float sum = sigmoidScore; + sum += __shfl_xor_sync(0xffffffff, sum, 4); + sum += __shfl_xor_sync(0xffffffff, sum, 2); + sum += __shfl_xor_sync(0xffffffff, sum, 1); + if (isTop) + { + auto const& post = params.mExpertSelectParams.mPostprocessParams; + float const weight = sigmoidScore * post.routeScale / (sum + post.sumEpsilon); + smemKIdx[myExpert * kKIdxRowBytes + warpIdx] = static_cast(laneIdx); + smemTopKExpert[warpIdx * kTopK + laneIdx] = static_cast(myExpert); + atomicOr(&smemSelected[myExpert / 32], 1u << (myExpert % 32)); + params.mPtrTopKWeights[warpIdx * kTopK + laneIdx] = __nv_bfloat16{weight}; + } + } + __syncthreads(); + + // ----- Metadata: tile index = rank among the selected local experts ----- + int32_t const localBegin = params.mLocalExpertsStartIdx; + int32_t const localEnd = params.mLocalExpertsStartIdx + params.mNumLocalExperts; + uint4 const lo = *reinterpret_cast(&smemSelected[0]); + uint4 const hi = *reinterpret_cast(&smemSelected[4]); + uint32_t selected[8] = {lo.x, lo.y, lo.z, lo.w, hi.x, hi.y, hi.z, hi.w}; + if (localBegin != 0 || localEnd < kNumExperts) + { +#pragma unroll + for (int i = 0; i < 8; ++i) + { + int32_t const b0 = max(localBegin - 32 * i, 0); + int32_t const b1 = min(localEnd - 32 * i, 32); + uint32_t const upTo = b1 >= 32 ? 0xFFFFFFFFu : ((1u << max(b1, 0)) - 1u); + uint32_t const below = b0 >= 32 ? 0xFFFFFFFFu : ((1u << b0) - 1u); + selected[i] &= b1 <= b0 ? 0u : (upTo & ~below); + } + } + int32_t numSelected = 0; +#pragma unroll + for (int i = 0; i < 8; ++i) + { + numSelected += __popc(selected[i]); + } + // (rank, is selected and local) of an expert, from masked popcounts of every word: indexing the words by the expert + // would place them in local memory. + auto rankOf = [&](int32_t e, bool& isSelected) + { + int32_t const word = e >> 5; + uint32_t const bit = 1u << (e & 31); + int32_t rank = 0; + int32_t hit = 0; +#pragma unroll + for (int i = 0; i < 8; ++i) + { + rank += __popc(selected[i] & (i < word ? 0xFFFFFFFFu : (i == word ? bit - 1u : 0u))); + hit += __popc(selected[i] & (i == word ? bit : 0u)); + } + isSelected = hit != 0; + return rank; + }; + auto tileRows = [&](int32_t n) + { + return params.mIsPow2 ? mulLog2(n, params.mPaddingLog2) : mulTileN(n, params.mTileTokensDim); + }; + + { + int32_t const expert = threadIdx.x; + bool isSelected; + int32_t const rank = rankOf(expert, isSelected); + if (isSelected) + { + uint2 const kRow = *reinterpret_cast(&smemKIdx[expert * kKIdxRowBytes]); + int32_t const count = __popc(__vsetne4(kRow.x, 0xFFFFFFFFu)) + __popc(__vsetne4(kRow.y, 0xFFFFFFFFu)); + params.mPtrCtaIdxXyToBatchIdx[rank] = expert - localBegin; + params.mPtrCtaIdxXyToMnLimit[rank] = min(tileRows(rank + 1), tileRows(rank) + count); + } + } + if (threadIdx.x == 0) + { + params.mPtrPermutedIdxSize[0] = tileRows(numSelected); + params.mPtrNumNonExitingCtas[0] = numSelected; + } + + // Permutation, one (token, k) pair per thread. offset-within-expert = number of earlier tokens that also selected + // the expert (its byte-map row below byte t), the small-token kernel's order. + if (static_cast(threadIdx.x) < kNumTokens * kTopK) + { + int32_t const expandedIdx = threadIdx.x; + int32_t const tokenIdx = expandedIdx / kTopK; + int32_t const pairExpert = smemTopKExpert[expandedIdx]; + bool pairIsLocal; + int32_t const rank = rankOf(pairExpert, pairIsLocal); + uint2 const row = *reinterpret_cast(&smemKIdx[pairExpert * kKIdxRowBytes]); + uint64_t const present = (static_cast(__vsetne4(row.y, 0xFFFFFFFFu)) << 32) + | static_cast(__vsetne4(row.x, 0xFFFFFFFFu)); + uint64_t const earlier = tokenIdx == 0 ? 0ull : (present & ((1ull << (8 * tokenIdx)) - 1)); + int32_t const permutedIdx = pairIsLocal ? tileRows(rank) + __popcll(earlier) : int32_t{-1}; + if (params.mPtrExpandedIdxToPermutedIdx != nullptr) + { + params.mPtrExpandedIdxToPermutedIdx[expandedIdx] = permutedIdx; + } + if (params.mPtrPermutedIdxToExpandedIdx != nullptr && pairIsLocal) + { + params.mPtrPermutedIdxToExpandedIdx[permutedIdx] = expandedIdx; + } + if (params.mPtrPermutedIdxToTokenIdx != nullptr && pairIsLocal) + { + params.mPtrPermutedIdxToTokenIdx[permutedIdx] = tokenIdx; + } + } + if (threadIdx.x == 0) + { + *epochPtr = epoch + 1u; + } + // After all routing outputs are written. + cudaTriggerProgrammaticLaunchCompletion(); +} +#else +__global__ void __launch_bounds__(kNumThreads, 1) routingFusedRouterKernel( + FusedRouterParams const __grid_constant__ /* p */) +{ + assert(false && "routingFusedRouterKernel is only supported on SM100+ architectures"); +} +#endif // if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) + +template +void launchFusedRouter(FusedRouterParams const& params, bool usePdl, cudaStream_t stream) +{ + cudaLaunchConfig_t config{}; + config.gridDim = kNumGemvCtas + 1; + config.blockDim = kNumThreads; + config.dynamicSmemBytes = kSmemBytes; + config.stream = stream; + cudaLaunchAttribute attributes[1] = {}; + attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attributes[0].val.programmaticStreamSerializationAllowed = int(usePdl); + config.attrs = attributes; + config.numAttrs = 1; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, routingFusedRouterKernel, params)); +} + +} // namespace + +bool supportsFusedRouter(Data const& data, int32_t hiddenDim) +{ + return hiddenDim == kHiddenDim && data.mNumExperts == kNumExperts && data.mTopK == kTopK + && data.mNumTokens >= 1 && data.mNumTokens <= kMaxTokens && data.mNumTokens <= data.mTileTokensDim + && data.mLocalExpertsStrideLog2 == 0 && data.mPreprocessType == RoutingPreprocessType::SigmoidBias + && data.mPostprocessType == RoutingPostprocessType::ScaledSumNormalize + && data.mDtypeOutput == tg::Dtype::Bfloat16 && data.mPtrTopKWeights != nullptr + && data.mPtrRoutingBias != nullptr && data.mPtrTopKIds == nullptr; +} + +void runFusedRouter(Data const& data, void const* hiddenStates, void const* routerWeight, int32_t hiddenDim, + int32_t* workspace, void* stream) +{ + TLLM_CHECK_WITH_INFO(supportsFusedRouter(data, hiddenDim), + "The fused router + routing kernel supports 1-8 tokens (at most the tile size), 256 experts, top-8, hidden " + "size %d and SigmoidBias + ScaledSumNormalize routing; got %d tokens, tile %d, %d experts, top-%d, hidden %d", + kHiddenDim, data.mNumTokens, data.mTileTokensDim, data.mNumExperts, data.mTopK, hiddenDim); + FusedRouterParams params; + params.routing = FusedKernelParams::setKernelParams(data); + params.hiddenStates = static_cast<__nv_bfloat16 const*>(hiddenStates); + params.routerWeight = static_cast<__nv_bfloat16 const*>(routerWeight); + params.workspace = workspace; + auto const s = static_cast(stream); + switch (data.mNumTokens) + { + case 1: launchFusedRouter<1>(params, data.mUsePdl, s); break; + case 2: launchFusedRouter<2>(params, data.mUsePdl, s); break; + case 3: launchFusedRouter<3>(params, data.mUsePdl, s); break; + case 4: launchFusedRouter<4>(params, data.mUsePdl, s); break; + case 5: launchFusedRouter<5>(params, data.mUsePdl, s); break; + case 6: launchFusedRouter<6>(params, data.mUsePdl, s); break; + case 7: launchFusedRouter<7>(params, data.mUsePdl, s); break; + default: launchFusedRouter<8>(params, data.mUsePdl, s); break; + } +} + +} // namespace routingCustom +} // namespace moe::dev::routing diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingKernel.h b/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingKernel.h index 3f238844e7fc..668aeb8b9ec2 100644 --- a/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingKernel.h +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingKernel.h @@ -384,6 +384,21 @@ void launchBlockKernel(Data const& data, uint32_t numThreadsHist, void* stream); void launchCoopBlockKernel(Data const& data, uint32_t numThreadsHist, void* stream); void launchSmallTokenBlockKernel(Data const& data, uint32_t numThreadsHist, void* stream); +// Workspace of runFusedRouter, in 32-bit words: two exchange buffers of [8 tokens][256 experts][4 slices] fp32 whose +// idle value is -0.0 (bits 0x80000000), then the epoch (0 before the first launch), padded to 16 bytes. +static constexpr int32_t FusedRouterWorkspaceWords = 2 * 8 * 256 * 4 + 4; + +// Whether runFusedRouter covers this routing configuration and hidden size: 1-8 tokens and at most mTileTokensDim, +// 256 experts, top-8, unit local-expert stride, SigmoidBias + ScaledSumNormalize, bf16 expert weights. +bool supportsFusedRouter(Data const& data, int32_t hiddenDim); + +// The router GEMV of dsv3_router_gemm_op (bf16 hiddenStates [mNumTokens, hiddenDim] x routerWeight [mNumExperts, +// hiddenDim], the same fp32 logits, which are not stored) and the routing of routingIndicesSmallTokenKernel in one +// kernel, with the same routing outputs; mPtrScores is not read. workspace: FusedRouterWorkspaceWords words, owned by +// the caller and initialized once; launches that share a workspace must run in stream order. +void runFusedRouter(Data const& data, void const* hiddenStates, void const* routerWeight, int32_t hiddenDim, + int32_t* workspace, void* stream); + } // namespace routingCustom //////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.cu b/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.cu index ab5c9e5d945d..4fe30ea94507 100644 --- a/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.cu +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.cu @@ -411,6 +411,44 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3 serializeMoeRoutingMethodType(routingMethodType).c_str(), (int) routingMethodType); } } + +void Runner::runFusedRouter(void const* routerHiddenStates, void const* routerWeight, int32_t hiddenDim, + int32_t* routerWorkspace, void* routingBias, int32_t numTokens, int32_t numExperts, int32_t topK, + int32_t localExpertOffset, int32_t localNumExperts, float routedScalingFactor, int32_t* permutedIdxSize, + int32_t* expandedIdxToPermutedIdx, int32_t* permutedIdxToExpandedIdx, int32_t* permutedIdxToTokenIdx, + void* expertWeights, int32_t* ctaIdxXyToBatchIdx, int32_t* ctaIdxXyToMnLimit, int32_t* numNonExitingCtas, + cudaStream_t stream, btg::Dtype dtypeRoutingBias) +{ + // The routing configuration of run()'s DeepSeekV3 branch without groups. + moe::dev::routing::routingCustom::Data routingData; + routingData.mDtypeOutput = btg::Dtype::Bfloat16; + routingData.mDtypeInput = btg::Dtype::Fp32; + routingData.mUsePdl = tensorrt_llm::common::getEnvEnablePDL(); + routingData.mPreprocessType = moe::dev::routing::RoutingPreprocessType::SigmoidBias; + routingData.mPostprocessType = moe::dev::routing::RoutingPostprocessType::ScaledSumNormalize; + routingData.mPtrRoutingBias = routingBias; + routingData.mDtypeBias = dtypeRoutingBias; + routingData.mRouteScale = routedScalingFactor; + routingData.mPtrPermutedIdxSize = permutedIdxSize; + routingData.mPtrExpandedIdxToPermutedIdx = expandedIdxToPermutedIdx; + routingData.mPtrPermutedIdxToExpandedIdx = permutedIdxToExpandedIdx; + routingData.mPtrPermutedIdxToTokenIdx = permutedIdxToTokenIdx; + routingData.mPtrTopKWeights = expertWeights; + routingData.mPtrCtaIdxXyToBatchIdx = ctaIdxXyToBatchIdx; + routingData.mPtrCtaIdxXyToMnLimit = ctaIdxXyToMnLimit; + routingData.mPtrNumNonExitingCtas = numNonExitingCtas; + routingData.mNumTokens = numTokens; + routingData.mNumExperts = numExperts; + routingData.mTopK = topK; + routingData.mPaddingLog2 = computeLog2(mTileTokensDim); + routingData.mTileTokensDim = mTileTokensDim; + routingData.mLocalExpertsStartIdx = localExpertOffset; + routingData.mLocalExpertsStrideLog2 = 0; + routingData.mNumLocalExperts = localNumExperts; + + moe::dev::routing::routingCustom::runFusedRouter( + routingData, routerHiddenStates, routerWeight, hiddenDim, routerWorkspace, stream); +} } // namespace Routing namespace PermuteGemm1 diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.h b/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.h index 8f7ee3a84979..cbe79c27e961 100644 --- a/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.h +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.h @@ -193,6 +193,17 @@ class Runner batchedGemm::trtllm::gen::Dtype dtypeRoutingLogits = batchedGemm::trtllm::gen::Dtype::Bfloat16, batchedGemm::trtllm::gen::Dtype dtypeRoutingBias = batchedGemm::trtllm::gen::Dtype::Bfloat16); + // run() for DeepSeekV3 routing without groups or fused shared experts, with the router GEMV in the same kernel: + // the routing logits are computed from routerHiddenStates [numTokens, hiddenDim] and routerWeight [numExperts, + // hiddenDim] (bf16) as dsv3_router_gemm_op computes them (see routingCustom::runFusedRouter, which checks that the + // shape is covered). + void runFusedRouter(void const* routerHiddenStates, void const* routerWeight, int32_t hiddenDim, + int32_t* routerWorkspace, void* routingBias, int32_t numTokens, int32_t numExperts, int32_t topK, + int32_t localExpertOffset, int32_t localNumExperts, float routedScalingFactor, int32_t* permutedIdxSize, + int32_t* expandedIdxToPermutedIdx, int32_t* permutedIdxToExpandedIdx, int32_t* permutedIdxToTokenIdx, + void* expertWeights, int32_t* ctaIdxXyToBatchIdx, int32_t* ctaIdxXyToMnLimit, int32_t* numNonExitingCtas, + cudaStream_t stream, batchedGemm::trtllm::gen::Dtype dtypeRoutingBias); + private: int32_t mTileTokensDim; }; diff --git a/cpp/tensorrt_llm/thop/moe/trtllmGen/fp4BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/moe/trtllmGen/fp4BlockScaleMoe.cpp index d054b29b273a..4b0859f02175 100644 --- a/cpp/tensorrt_llm/thop/moe/trtllmGen/fp4BlockScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/moe/trtllmGen/fp4BlockScaleMoe.cpp @@ -47,7 +47,10 @@ std::vector run_fp4_block_scale_moe_runner(torch::optional const& topk_weights, torch::optional const& topk_ids, torch::optional const& out_tensor = torch::nullopt, - std::optional const hidden_states_ready_event = std::nullopt) + std::optional const hidden_states_ready_event = std::nullopt, + torch::optional const& router_hidden_states = torch::nullopt, + torch::optional const& router_weight = torch::nullopt, + torch::optional const& router_workspace = torch::nullopt) { TORCH_CHECK(dtype == btg::Dtype::E4m3 || dtype == btg::Dtype::E2m1, "dtype can only be e4m3 or e2m1."); TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "Only SM100f is supported by FP4 block scale MOE"); @@ -291,16 +294,50 @@ std::vector run_fp4_block_scale_moe_runner(torch::optional(), - expert_count_histogram.data_ptr(), total_num_padded_tokens.data_ptr(), - expanded_idx_to_permuted_idx.data_ptr(), nullptr, /*permuted_idx_to_expanded_idx.data_ptr(),*/ - permuted_idx_to_token_idx.data_ptr(), expert_weights_ptr, args.topk_ids, - num_tokens_per_expert.data_ptr(), cta_idx_xy_to_batch_idx.data_ptr(), - cta_idx_xy_to_mn_limit.data_ptr(), num_non_exiting_ctas.data_ptr(), args.mDtypeElt, - false /* use_routing_scales_on_input */, false /* use_deep_seek_fp8 */, - static_cast(routing_method_type), stream, dtypeRoutingLogits, args.mDtypeBias); + if (router_hidden_states.has_value()) + { + // The router GEMV runs in the routing kernel: routing_logits only gives the shape and is not read or written. + TORCH_CHECK(router_weight.has_value() && router_workspace.has_value(), + "router_hidden_states needs router_weight and router_workspace."); + TORCH_CHECK(!args.topk_ids + && static_cast(routing_method_type) == RoutingMethodType::DeepSeekV3 + && args.n_group <= 1 && routing_bias.has_value(), + "The fused router covers DeepSeekV3 routing without groups, with a routing bias, from the router inputs."); + auto const& hs = router_hidden_states.value(); + auto const& rw = router_weight.value(); + auto const& ws = router_workspace.value(); + TORCH_CHECK(hs.scalar_type() == at::ScalarType::BFloat16 && hs.dim() == 2 && hs.is_contiguous() + && hs.sizes()[0] == args.num_tokens && hs.sizes()[1] == args.hidden_size + && hs.device() == routing_device, + "router_hidden_states must be a contiguous bf16 [num_tokens, hidden_size] tensor on the routing device."); + TORCH_CHECK(rw.scalar_type() == at::ScalarType::BFloat16 && rw.dim() == 2 && rw.is_contiguous() + && rw.sizes()[0] == num_experts && rw.sizes()[1] == hs.sizes()[1] && rw.device() == routing_device, + "router_weight must be a contiguous bf16 [num_experts, hidden] tensor on the routing device."); + TORCH_CHECK(ws.scalar_type() == at::ScalarType::Int && ws.is_contiguous() + && ws.numel() == moe::dev::routing::routingCustom::FusedRouterWorkspaceWords + && ws.device() == routing_device, + "router_workspace must be a contiguous int32 tensor of ", + moe::dev::routing::routingCustom::FusedRouterWorkspaceWords, " words on the routing device."); + routing_runner.runFusedRouter(hs.data_ptr(), rw.data_ptr(), static_cast(hs.sizes()[1]), + ws.data_ptr(), args.routing_bias, args.num_tokens, args.num_experts, args.top_k, + args.local_expert_offset, args.local_num_experts, args.routed_scaling_factor, + total_num_padded_tokens.data_ptr(), expanded_idx_to_permuted_idx.data_ptr(), nullptr, + permuted_idx_to_token_idx.data_ptr(), expert_weights_ptr, cta_idx_xy_to_batch_idx.data_ptr(), + cta_idx_xy_to_mn_limit.data_ptr(), num_non_exiting_ctas.data_ptr(), stream, args.mDtypeBias); + } + else + { + routing_runner.run(args.routing_logits, args.routing_bias, args.num_tokens, args.num_experts, args.top_k, + /* num_fused_shared_expert */ 0, args.n_group, args.topk_group, args.local_expert_offset, + args.local_num_experts, args.routed_scaling_factor, expert_indexes.data_ptr(), + expert_count_histogram.data_ptr(), total_num_padded_tokens.data_ptr(), + expanded_idx_to_permuted_idx.data_ptr(), nullptr, /*permuted_idx_to_expanded_idx.data_ptr(),*/ + permuted_idx_to_token_idx.data_ptr(), expert_weights_ptr, args.topk_ids, + num_tokens_per_expert.data_ptr(), cta_idx_xy_to_batch_idx.data_ptr(), + cta_idx_xy_to_mn_limit.data_ptr(), num_non_exiting_ctas.data_ptr(), args.mDtypeElt, + false /* use_routing_scales_on_input */, false /* use_deep_seek_fp8 */, + static_cast(routing_method_type), stream, dtypeRoutingLogits, args.mDtypeBias); + } // // FC13 (gemm1) + FC2 (gemm2) @@ -595,7 +632,10 @@ class FP4BlockScaleMoeRunner : public torch::CustomClassHolder std::optional const routed_scaling_factor, int64_t const routing_method_type, bool const do_finalize, std::vector moeConfigIndex, torch::optional const& topk_weights, torch::optional const& topk_ids, torch::optional const& output = torch::nullopt, - std::optional const hidden_states_ready_event = std::nullopt) + std::optional const hidden_states_ready_event = std::nullopt, + torch::optional const& router_hidden_states = torch::nullopt, + torch::optional const& router_weight = torch::nullopt, + torch::optional const& router_workspace = torch::nullopt) { int64_t tileN = moeConfigIndex[0]; int64_t config = moeConfigIndex[1]; @@ -645,7 +685,8 @@ class FP4BlockScaleMoeRunner : public torch::CustomClassHolder gemm2_weights_scale, gemm2_bias, output1_scales_scalar, output1_scales_gate_scalar, output2_scales_scalar, num_experts, top_k, n_group, topk_group, intermediate_size, local_expert_offset, local_num_experts, routed_scaling_factor, tileN, routing_method_type, do_finalize, mDtypeElt, *runner, configToRun, - topk_weights, topk_ids, output, hidden_states_ready_event); + topk_weights, topk_ids, output, hidden_states_ready_event, router_hidden_states, router_weight, + router_workspace); } private: diff --git a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py index 053d55fc50f1..6e43eacb9078 100644 --- a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py @@ -500,6 +500,9 @@ def forward( tactic: List[int] = [-1, -1], output: Optional[torch.Tensor] = None, hidden_states_ready_event: Optional[int] = None, + router_hidden_states: Optional[torch.Tensor] = None, + router_weight: Optional[torch.Tensor] = None, + router_workspace: Optional[torch.Tensor] = None, ) -> torch.Tensor: assert isinstance(tactic, list) @@ -518,7 +521,8 @@ def forward( self.local_expert_offset, self.local_num_experts, self.routed_scaling_factor, self.routing_method_type, self.do_finalize, tactic, args.topk_weights, args.topk_ids, output, - hidden_states_ready_event) + hidden_states_ready_event, router_hidden_states, router_weight, + router_workspace) def get_valid_tactics(self, inputs: List[torch.Tensor], profile: OptimizationProfile, @@ -682,11 +686,24 @@ def fp4_block_scale_moe_runner(routing_logits: Optional[torch.Tensor], output: Optional[torch.Tensor] = None, tune_max_num_tokens: int = 8192, use_dp: bool = False, - hidden_states_ready_event: Optional[int] = None + hidden_states_ready_event: Optional[int] = None, + router_hidden_states: Optional[ + torch.Tensor] = None, + router_weight: Optional[torch.Tensor] = None, + router_workspace: Optional[torch.Tensor] = None ) -> List[torch.Tensor]: """hidden_states_ready_event: raw cudaEvent_t (torch.cuda.Event.cuda_event) recorded on another stream once hidden_states and hidden_states_scale are - written; the op waits on it right before FC1, after launching routing.""" + written; the op waits on it right before FC1, after launching routing. + + router_hidden_states, router_weight, router_workspace: the router GEMV runs + in the routing kernel (DeepSeekV3 routing without groups, 256 experts, + top-8, hidden 6144, 1-8 tokens): the logits of + dsv3_router_gemm_op(router_hidden_states, router_weight.t()) are computed + there and routing_logits only gives their shape (it is not read or + written). router_workspace: int32 [FUSED_ROUTER_WORKSPACE_WORDS], see + make_fused_router_workspace; the kernel uses and restores it. Profiling + runs of the autotuner route from routing_logits as usual.""" tuner = AutoTuner.get() kernel_runner = FP4BlockScaleMoERunner( @@ -763,7 +780,10 @@ def fp4_block_scale_moe_runner(routing_logits: Optional[torch.Tensor], input_tensors, tactic=[-1, -1] if best_tactic == -1 else best_tactic, output=output, - hidden_states_ready_event=hidden_states_ready_event) + hidden_states_ready_event=hidden_states_ready_event, + router_hidden_states=router_hidden_states, + router_weight=router_weight, + router_workspace=router_workspace) # When output is provided and do_finalize=True, the result is written in-place to output. # Return empty tensor to avoid aliasing constraint violation in PyTorch 2.9.1+ # (custom op output cannot be the same tensor as input). @@ -773,6 +793,24 @@ def fp4_block_scale_moe_runner(routing_logits: Optional[torch.Tensor], return result +# routingCustom::FusedRouterWorkspaceWords: two exchange buffers of +# [8 tokens][256 experts][4 slices] fp32, then the epoch word, padded to 16 B. +FUSED_ROUTER_WORKSPACE_WORDS = 2 * 8 * 256 * 4 + 4 + + +def make_fused_router_workspace(device: torch.device) -> torch.Tensor: + """Workspace of fp4_block_scale_moe_runner's fused router GEMV + routing + kernel (router_workspace): both exchange buffers idle (-0.0, bits + 0x80000000) and epoch 0. Between launches the buffer the next launch uses + is idle; nothing else may write the workspace. Launches that share a + workspace must run in stream order.""" + workspace = torch.zeros(FUSED_ROUTER_WORKSPACE_WORDS, + dtype=torch.int32, + device=device) + workspace[:FUSED_ROUTER_WORKSPACE_WORDS - 4].fill_(-2**31) + return workspace + + def fp4_block_scale_fake_output_without_finalize( hidden_states: Union[torch.Tensor, Fp4QuantizedTensor], num_experts: int, @@ -832,7 +870,10 @@ def _(routing_logits, output: Optional[torch.Tensor] = None, tune_max_num_tokens: int = 8192, use_dp: bool = False, - hidden_states_ready_event: Optional[int] = None) -> List[torch.Tensor]: + hidden_states_ready_event: Optional[int] = None, + router_hidden_states: Optional[torch.Tensor] = None, + router_weight: Optional[torch.Tensor] = None, + router_workspace: Optional[torch.Tensor] = None) -> List[torch.Tensor]: if do_finalize: num_tokens = hidden_states.shape[0] hidden_size = hidden_states.shape[1] * 2 From 7ef4ef669178872a6b032ed45c86777f06857a79 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:04:49 -0700 Subject: [PATCH 16/53] Add a K = 128 FC2 kernel to the TRTLLM-Gen NVFP4 MoE With prefer_gemm2_k128 the NVFP4 MoE runs FC2 on moe::dev::gemm2K128 where it covers the call: E2M1 x E2M1 without DeepSeek FP8 or fine-grained sync, tile N 8, intermediate size 128 without padding, no bias or fused shared experts, at most 64 token tiles and at most 15 units (token tile x 128 output features) per SM. At that K the Gemm2 cubins spend half of their 256-wide K tile on zero padding. gemm2K128 runs one persistent CTA per SM over a contiguous run of units: weights and their block scales by TMA before the grid-dependency wait, the FC1 output after it, two tcgen05 mxf4nvf4 block16 MMAs (M 128, N 8, K 64) per unit, then the cubin's FP32 output scale and BF16 rounding, so the rows are the cubin's bitwise. It reads the routing outputs before its wait (its primary, FC1, triggers only after its own wait on routing), triggers its dependents right after the wait and prefers the maximum shared-memory carveout. Everywhere else, and without the flag, the Gemm2 cubin runs; the autotuner keys a flagged runner apart. Signed-off-by: Dongfeng Yu --- .../kernels/moe/trtllmGen/Gemm2K128.cu | 667 ++++++++++++++++++ .../kernels/moe/trtllmGen/Gemm2K128.h | 60 ++ .../kernels/moe/trtllmGen/runner.cu | 56 +- .../kernels/moe/trtllmGen/runner.h | 9 + .../thop/moe/trtllmGen/fp4BlockScaleMoe.cpp | 7 +- .../custom_ops/trtllm_gen_custom_ops.py | 26 +- 6 files changed, 808 insertions(+), 17 deletions(-) create mode 100644 cpp/tensorrt_llm/kernels/moe/trtllmGen/Gemm2K128.cu create mode 100644 cpp/tensorrt_llm/kernels/moe/trtllmGen/Gemm2K128.h diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/Gemm2K128.cu b/cpp/tensorrt_llm/kernels/moe/trtllmGen/Gemm2K128.cu new file mode 100644 index 000000000000..651cb992ec02 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/Gemm2K128.cu @@ -0,0 +1,667 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// FC2 of the NVFP4 MoE at intermediate size 128 (K = 128) and tile N 8 (see Gemm2K128.h for the contract). +// +// Work. A unit is (token tile t, 128-row tile of the output features); the units are split evenly over a persistent +// grid of one CTA per SM, each CTA owning a contiguous run of units (one or two token tiles). A CTA holds every unit of +// its run at once: the weight tiles (8 KB of E2M1, SWIZZLE_64B, by TMA in two-tile boxes) and their 128x4 block scales +// (one bulk copy per token tile) are requested before the grid-dependency wait, the activations of its token tiles +// (8 rows x 64 B) and their 8x4 scales after it. +// +// MMA. tcgen05.mma kind::mxf4nvf4 block16, M = 128 (weight rows), N = 8 (token rows), two K = 64 steps per unit, issued +// by the elected lane of the converged MMA warp, one commit per unit; the weight scales go into TMEM by tcgen05.cp +// ahead of the MMAs (a separate warp, as their bulk copy lands), the token scales after the wait. The TRTLLM-Gen cubin +// computes the same two K steps and two more on the zero padding of its 256-wide K tile; adding +0.0f to the +// accumulator gives its result, bitwise. +// +// Accumulators. Unit k of a token tile's run writes its 8 TMEM columns at base + stride * k, stride = the tile's valid +// rows rounded up to 2, 4 or 8 (accumulators are 2-column aligned): the valid rows of all units of the tile lie in +// stride * units consecutive columns, and the tail columns a unit shares with the next one are padding rows, written +// in issue order. The epilogue reads them in 16-column loads, each once its last unit is complete. +// +// Warps: 0-3 epilogue (TMEM lanes 32 w .. 32 w + 31 = weight rows of each tile), 4 weight loads, 5 TMEM allocation, +// activations and MMA, 6 weight-scale copies. + +#include "Gemm2K128.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaDriverWrapper.h" +#include "tensorrt_llm/common/cudaUtils.h" + +#include +#include + +namespace moe::dev::gemm2K128 +{ +namespace +{ + +constexpr int kTileRows = 128; +constexpr int kK = 128; +constexpr int kTokens = 8; +constexpr int kTileBytesA = kTileRows * kK / 2; // 8,192: 128 rows x 64 B, SWIZZLE_64B +constexpr int kTileBytesSf = kTileRows * kK / 16; // 1,024: two 512 B blocks of the 128x4 layout +constexpr int kActBytes = kTokens * kK / 2; // 512: one SWIZZLE_64B atom +constexpr int kSfbBytes = 512; // 32 rows x 16 B: the tcgen05.cp source of a token tile's scales +constexpr int kMaxUnits = 15; // units per CTA +constexpr int kMaxSlots = 2; // token tiles per CTA (a run of kMaxUnits units spans at most two) +constexpr int kMaxTokenTiles = 64; +constexpr int kThreads = 224; +constexpr uint32_t kTmemCols = 512; +constexpr uint32_t kTmemAcc = 0; // packed accumulators, at most 256 columns +constexpr uint32_t kTmemSfa = 256; // kMaxUnits x 8 columns: the weight scales of every unit +constexpr uint32_t kTmemSfb = 376; // kMaxSlots x 4 columns: per K step a column of token scales and a zero column +// Instruction descriptor, kind::mxf4nvf4: A/B E2M1, UE4M3 scales, K-major, M = 128, N = 8. +constexpr uint32_t kIdesc = (1u << 7) | (1u << 10) | ((8u >> 3) << 17) | ((128u >> 4) << 24); +// K-major SWIZZLE_64B operand descriptor without its start address: 8-row x 64 B atoms 512 B apart, version 1. +constexpr uint64_t kDescSw64 = (1ull << 16) | (static_cast(512 >> 4) << 32) | (1ull << 46) | (4ull << 61); +// tcgen05.cp 32x128b source descriptor without its start address: 8-row x 16 B core matrices 128 B apart. +constexpr uint64_t kDescSf = (static_cast(128 >> 4) << 32) | (1ull << 46); +constexpr int kSmemBytes = kMaxUnits * (kTileBytesA + kTileBytesSf) + kMaxSlots * (kActBytes + kSfbBytes) + 1024; + +struct Params +{ + CUtensorMap tmW2; // weights [experts * hidden rows][64 B], box {64, 256}, SWIZZLE_64B + CUtensorMap tmW1; // the same tensor, box {64, 128} + uint8_t const* wSf; + uint8_t const* act; + uint8_t const* actSf; + float const* scale; + int32_t const* numNonExitingCtas; + int32_t const* ctaIdxXyToBatchIdx; + int32_t const* ctaIdxXyToMnLimit; + __nv_bfloat16* out; + int32_t hidden; + int32_t maxCtas; +}; + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1100 + +__device__ __forceinline__ uint32_t smemU32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ void mbarInit(uint64_t* bar, uint32_t count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" ::"r"(smemU32(bar)), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarExpectTx(uint64_t* bar, uint32_t bytes) +{ + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(smemU32(bar)), "r"(bytes) : "memory"); +} + +__device__ __forceinline__ void mbarWait(uint64_t* bar, uint32_t parity) +{ + uint32_t ok = 0; + while (!ok) + { + asm volatile( + "{\n" + ".reg .pred p;\n" + "mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n" + "selp.u32 %0, 1, 0, p;\n" + "}\n" + : "=r"(ok) + : "r"(smemU32(bar)), "r"(parity) + : "memory"); + } +} + +__device__ __forceinline__ void tmaLoad2d(void* dst, void const* map, int x, int y, uint64_t* bar) +{ + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1, {%2, %3}], [%4];" ::"r"( + smemU32(dst)), + "l"(map), "r"(x), "r"(y), "r"(smemU32(bar)) + : "memory"); +} + +__device__ __forceinline__ void bulkLoad(void* dst, void const* src, uint32_t bytes, uint64_t* bar) +{ + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( + smemU32(dst)), + "l"(src), "r"(bytes), "r"(smemU32(bar)) + : "memory"); +} + +__device__ __forceinline__ int ldRelaxed(int32_t const* p) +{ + int v; + asm volatile("ld.relaxed.gpu.global.b32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ void namedBarSync(int id, int threads) +{ + asm volatile("bar.sync %0, %1;" ::"r"(id), "r"(threads) : "memory"); +} + +__device__ __forceinline__ void namedBarArrive(int id, int threads) +{ + asm volatile("bar.arrive %0, %1;" ::"r"(id), "r"(threads) : "memory"); +} + +__device__ __forceinline__ void fenceProxyAsyncSmem() +{ + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + +__device__ __forceinline__ void tcFenceBefore() +{ + asm volatile("tcgen05.fence::before_thread_sync;" ::: "memory"); +} + +__device__ __forceinline__ void tcFenceAfter() +{ + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); +} + +__device__ __forceinline__ void tmemAlloc(uint32_t* dst, uint32_t cols) +{ + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"(smemU32(dst)), "r"(cols) + : "memory"); +} + +__device__ __forceinline__ void tmemRelinquish() +{ + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;" ::: "memory"); +} + +__device__ __forceinline__ void tmemDealloc(uint32_t taddr, uint32_t cols) +{ + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(taddr), "r"(cols) : "memory"); +} + +// The following three are executed by a whole (converged) warp; its elected lane issues the instruction. +__device__ __forceinline__ void utccpElect(uint32_t taddr, uint64_t sdesc) +{ + asm volatile( + "{\n" + ".reg .pred e;\n" + "elect.sync _|e, 0xffffffff;\n" + "@e tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;\n" + "}\n" ::"r"(taddr), + "l"(sdesc) + : "memory"); +} + +__device__ __forceinline__ void mmaElect(uint32_t d, uint64_t a, uint64_t b, uint32_t sfa, uint32_t sfb, uint32_t acc) +{ + asm volatile( + "{\n" + ".reg .pred e, p;\n" + "elect.sync _|e, 0xffffffff;\n" + "setp.ne.b32 p, %6, 0;\n" + "@e tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%4], [%5], p;\n" + "}\n" ::"r"(d), + "l"(a), "l"(b), "r"(kIdesc), "r"(sfa), "r"(sfb), "r"(acc) + : "memory"); +} + +__device__ __forceinline__ void commitElect(uint64_t* bar) +{ + asm volatile( + "{\n" + ".reg .pred e;\n" + "elect.sync _|e, 0xffffffff;\n" + "@e tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];\n" + "}\n" ::"r"(smemU32(bar)) + : "memory"); +} + +// Columns [taddr, taddr + 16) of the warp's TMEM lane quarter. +__device__ __forceinline__ void tmemLd16(uint32_t taddr, float (&v)[16]) +{ + uint32_t r[16]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32 {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, " + "%15}, [%16];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), "=r"(r[8]), + "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]) + : "r"(taddr) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" + : "+r"(r[0]), "+r"(r[1]), "+r"(r[2]), "+r"(r[3]), "+r"(r[4]), "+r"(r[5]), "+r"(r[6]), "+r"(r[7]), + "+r"(r[8]), "+r"(r[9]), "+r"(r[10]), "+r"(r[11]), "+r"(r[12]), "+r"(r[13]), "+r"(r[14]), "+r"(r[15]) + : + : "memory"); +#pragma unroll + for (int i = 0; i < 16; ++i) + { + v[i] = __uint_as_float(r[i]); + } +} + +// Stores the valid rows of the units [kBase, kBase + 16 / kStride) of a token tile from 16 packed accumulator columns +// (column kStride * k + r = row r of unit kBase + k). rowPtr: row 0 of the tile at this thread's feature of the run's +// first unit. +template +__device__ __forceinline__ void storeChunk( + float const (&v)[16], __nv_bfloat16* rowPtr, size_t rowStride, int kBase, int runLen, int nv, float scale) +{ + constexpr int kUnits = 16 / kStride; +#pragma unroll + for (int k = 0; k < kUnits; ++k) + { + if (kBase + k < runLen) + { + __nv_bfloat16* o = rowPtr + (kBase + k) * kTileRows; +#pragma unroll + for (int r = 0; r < kStride; ++r) + { + if (r < nv) + { + o[r * rowStride] = __float2bfloat16_rn(__fadd_rn(v[kStride * k + r], 0.0f) * scale); + } + } + } + } +} + +__device__ __forceinline__ int unitStart(int c, int units, int grid) +{ + return static_cast(static_cast(c) * units / grid); +} + +#endif + +__global__ void __launch_bounds__(kThreads, 1) gemm2K128Kernel(const __grid_constant__ Params p) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1100 + extern __shared__ __align__(1024) uint8_t smemRaw[]; + __shared__ __align__(8) uint64_t fullBar[kMaxUnits]; + __shared__ __align__(8) uint64_t sfaBar[kMaxSlots]; + __shared__ __align__(8) uint64_t sfaReady[kMaxUnits]; + __shared__ __align__(8) uint64_t accFull[kMaxUnits]; + __shared__ uint32_t sTmem; + __shared__ int sRange[2]; + __shared__ int sExpert[kMaxSlots]; + __shared__ int sNv[kMaxSlots]; + __shared__ int sStride[kMaxSlots]; + __shared__ int sCol[kMaxSlots]; + + uint8_t* smem = reinterpret_cast((reinterpret_cast(smemRaw) + 1023) & ~uintptr_t(1023)); + uint8_t* smemA = smem; + uint8_t* smemSf = smemA + kMaxUnits * kTileBytesA; + uint8_t* smemAct = smemSf + kMaxUnits * kTileBytesSf; + uint8_t* smemSfb = smemAct + kMaxSlots * kActBytes; + + int const warp = threadIdx.x >> 5; + int const lane = threadIdx.x & 31; + int const tilesPerSlot = p.hidden / kTileRows; + + if (warp == 4 && lane == 0) + { + for (int i = 0; i < kMaxUnits; ++i) + { + mbarInit(&fullBar[i], 1); + mbarInit(&sfaReady[i], 1); + mbarInit(&accFull[i], 1); + } + for (int j = 0; j < kMaxSlots; ++j) + { + mbarInit(&sfaBar[j], 1); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + asm volatile("prefetch.tensormap [%0];" ::"l"(&p.tmW2) : "memory"); + asm volatile("prefetch.tensormap [%0];" ::"l"(&p.tmW1) : "memory"); + } + if (warp == 5) + { + tmemAlloc(&sTmem, kTmemCols); + tmemRelinquish(); + } + tcFenceBefore(); + __syncthreads(); + tcFenceAfter(); + uint32_t const tmem = sTmem; + + if (warp == 4) + { + // The routing outputs are complete when this kernel starts: its primary, the FC1 cubin, triggers its dependents + // only after its own grid-dependency wait on the routing kernel. One round trip for the tile count and the + // experts and valid rows of the token tiles; the weights and scales are constants. + int const numTiles = ldRelaxed(p.numNonExitingCtas); + int const e0 = lane < p.maxCtas ? ldRelaxed(p.ctaIdxXyToBatchIdx + lane) : 0; + int const e1 = lane + 32 < p.maxCtas ? ldRelaxed(p.ctaIdxXyToBatchIdx + lane + 32) : 0; + int const m0 = lane < p.maxCtas ? ldRelaxed(p.ctaIdxXyToMnLimit + lane) : 0; + int const m1 = lane + 32 < p.maxCtas ? ldRelaxed(p.ctaIdxXyToMnLimit + lane + 32) : 0; + int const units = numTiles * tilesPerSlot; + int const a = unitStart(blockIdx.x, units, gridDim.x); + int const b = unitStart(blockIdx.x + 1, units, gridDim.x); + // Publish the run (its experts, valid rows and accumulator columns) before issuing the loads. + if (b > a) + { + int const tA = a / tilesPerSlot; + int const tB = (b - 1) / tilesPerSlot; + int col = 0; + for (int t = tA; t <= tB; ++t) + { + int const eLo = __shfl_sync(0xffffffffu, e0, t & 31); + int const eHi = __shfl_sync(0xffffffffu, e1, t & 31); + int const mLo = __shfl_sync(0xffffffffu, m0, t & 31); + int const mHi = __shfl_sync(0xffffffffu, m1, t & 31); + int const nv = (t < 32 ? mLo : mHi) - kTokens * t; + int const stride = nv <= 2 ? 2 : (nv <= 4 ? 4 : 8); + int const runLen = min(b, (t + 1) * tilesPerSlot) - max(a, t * tilesPerSlot); + if (lane == 0) + { + sExpert[t - tA] = t < 32 ? eLo : eHi; + sNv[t - tA] = nv; + sStride[t - tA] = stride; + sCol[t - tA] = col; + } + // The next token tile starts at a 32-column boundary (aligned TMEM loads). + col = (col + stride * runLen + 31) & ~31; + } + } + if (lane == 0) + { + sRange[0] = a; + sRange[1] = b; + } + __syncwarp(); + namedBarArrive(1, kThreads); + // Per token tile of the run: all its block scales in one bulk copy (first, for the scale-copy warp), then its + // weight tiles in two-tile boxes (a one-tile box for an odd last one). + if (lane == 0) + { + int i = 0; + int slot = 0; + while (a + i < b) + { + int const u0 = a + i; + int const t = u0 / tilesPerSlot; + int const nt0 = u0 - t * tilesPerSlot; + int const n = min(b - u0, tilesPerSlot - nt0); + mbarExpectTx(&sfaBar[slot], static_cast(n * kTileBytesSf)); + bulkLoad(smemSf + i * kTileBytesSf, + p.wSf + (static_cast(sExpert[slot]) * tilesPerSlot + nt0) * kTileBytesSf, + static_cast(n * kTileBytesSf), &sfaBar[slot]); + i += n; + ++slot; + } + i = 0; + slot = 0; + while (a + i < b) + { + int const u0 = a + i; + int const t = u0 / tilesPerSlot; + int const nt0 = u0 - t * tilesPerSlot; + int const n = min(b - u0, tilesPerSlot - nt0); + int const row0 = sExpert[slot] * p.hidden + nt0 * kTileRows; + for (int k = 0; k < n; k += 2) + { + bool const two = k + 1 < n; + mbarExpectTx(&fullBar[i + k], two ? 2 * kTileBytesA : kTileBytesA); + tmaLoad2d(smemA + (i + k) * kTileBytesA, two ? &p.tmW2 : &p.tmW1, 0, row0 + k * kTileRows, + &fullBar[i + k]); + } + i += n; + ++slot; + } + } + } + else + { + namedBarSync(1, kThreads); + } + int const a = sRange[0]; + int const b = sRange[1]; + int const n = b - a; + int const t0 = n > 0 ? a / tilesPerSlot : 0; + int const nt0 = a - t0 * tilesPerSlot; + int const numSlots = n > 0 ? (b - 1) / tilesPerSlot - t0 + 1 : 0; + + if (warp == 6) + { + // The weight scales of every unit into TMEM as their bulk copy lands (before the grid-dependency wait returns, + // for a CTA that runs while the previous kernel does); the MMA warp waits on sfaReady[unit]. + int u = 0; + int nt = nt0; + for (int j = 0; j < numSlots; ++j) + { + int const runLen = min(n - u, tilesPerSlot - nt); + mbarWait(&sfaBar[j], 0); + tcFenceAfter(); + for (int k = 0; k < runLen; ++k) + { + uint32_t const src = smemU32(smemSf + (u + k) * kTileBytesSf); + uint32_t const dst = tmem + kTmemSfa + 8 * (u + k); + utccpElect(dst, kDescSf | ((src & 0x3ffffu) >> 4)); + utccpElect(dst + 4, kDescSf | (((src + 512) & 0x3ffffu) >> 4)); + commitElect(&sfaReady[u + k]); + } + u += runLen; + nt = 0; + } + } + + // Everything below reads FC1's output or writes this kernel's output (the weight-load and scale-copy warps read + // neither). The dependents may launch once every CTA is past its wait. + if (warp < 4 || warp == 5) + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + if (threadIdx.x == 0) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + + if (warp < 4) + { + // Epilogue, per token tile of the run, in 16-column chunks, each once its last unit is complete. + uint32_t const tq = tmem + (static_cast(32 * warp) << 16); + // The weight rows are shuffled for the 128-row epilogue tile: row 32 w + l of a tile holds its feature + // 32 w + 4 (l & 7) + l / 8. + int const feature = 32 * warp + (lane & 7) * 4 + (lane >> 3); + int runStart = 0; + for (int j = 0; j < numSlots; ++j) + { + int const t = t0 + j; + int const ntFirst = j == 0 ? nt0 : 0; + int const runLen = min(n - runStart, tilesPerSlot - ntFirst); + int const nv = sNv[j]; + int const stride = sStride[j]; + float const scale = __ldg(p.scale + sExpert[j]); + int const unitsPerChunk = 16 / stride; + __nv_bfloat16* const rowBase + = p.out + static_cast(t) * kTokens * p.hidden + ntFirst * kTileRows + feature; + for (int kBase = 0; kBase < runLen; kBase += unitsPerChunk) + { + mbarWait(&accFull[runStart + min(kBase + unitsPerChunk, runLen) - 1], 0); + tcFenceAfter(); + float v[16]; + tmemLd16(tq + kTmemAcc + sCol[j] + stride * kBase, v); + if (stride == 2) + { + storeChunk<2>(v, rowBase, p.hidden, kBase, runLen, nv, scale); + } + else if (stride == 4) + { + storeChunk<4>(v, rowBase, p.hidden, kBase, runLen, nv, scale); + } + else + { + storeChunk<8>(v, rowBase, p.hidden, kBase, runLen, nv, scale); + } + } + runStart += runLen; + } + } + else if (warp == 5 && n > 0) + { + // Activations of the run's token tiles (8 rows of K = 128 E2M1, SWIZZLE_64B) and their 8x4 block scales, all + // loaded in one round trip. The scales go through shared memory (rows 0-7 of a 32-row block: per K step a + // column of 4 scales and a zero column) into every TMEM lane quarter. + int const r8 = lane >> 2; + int const cnk = lane & 3; + uint4 actV[kMaxSlots]; + uint32_t sf0[kMaxSlots]; + uint32_t sf1[kMaxSlots]; +#pragma unroll + for (int j = 0; j < kMaxSlots; ++j) + { + actV[j] = make_uint4(0, 0, 0, 0); + sf0[j] = 0; + sf1[j] = 0; + if (j < numSlots) + { + actV[j] = __ldcg(reinterpret_cast( + p.act + (static_cast(t0 + j) * kTokens + r8) * (kK / 2) + cnk * 16)); + if (lane < kTokens) + { + sf0[j] = __ldcg(reinterpret_cast(p.actSf) + 2 * (t0 + j) * kTokens + lane); + sf1[j] = __ldcg(reinterpret_cast(p.actSf) + (2 * (t0 + j) + 1) * kTokens + lane); + } + } + } +#pragma unroll + for (int j = 0; j < kMaxSlots; ++j) + { + if (j < numSlots) + { + *reinterpret_cast(smemAct + j * kActBytes + r8 * 64 + ((cnk ^ ((r8 >> 1) & 3)) << 4)) + = actV[j]; + *reinterpret_cast(smemSfb + j * kSfbBytes + lane * 16) = make_uint4(sf0[j], 0u, sf1[j], 0u); + } + } + fenceProxyAsyncSmem(); + __syncwarp(); + for (int j = 0; j < numSlots; ++j) + { + utccpElect(tmem + kTmemSfb + 4 * j, kDescSf | ((smemU32(smemSfb + j * kSfbBytes) & 0x3ffffu) >> 4)); + } + + int slot = 0; + int nt = nt0; + int runIdx = 0; // index of the unit in its token tile's run + uint32_t aBase = smemU32(smemA); + uint32_t bBase = smemU32(smemAct); + uint32_t sfb = tmem + kTmemSfb; + uint32_t d = tmem + kTmemAcc + sCol[0]; + uint32_t stride = sStride[0]; + for (int i = 0; i < n; ++i) + { + if (nt == tilesPerSlot) + { + ++slot; + nt = 0; + runIdx = 0; + bBase += kActBytes; + sfb += 4; + d = tmem + kTmemAcc + sCol[slot]; + stride = sStride[slot]; + } + if ((runIdx & 1) == 0) + { + // The first unit of a two-tile box. + mbarWait(&fullBar[i], 0); + } + mbarWait(&sfaReady[i], 0); + tcFenceAfter(); + uint32_t const sfa = tmem + kTmemSfa + 8 * i; + uint64_t const da = kDescSw64 | ((aBase & 0x3ffffu) >> 4); + uint64_t const db = kDescSw64 | ((bBase & 0x3ffffu) >> 4); + mmaElect(d, da, db, sfa, sfb, 0u); + mmaElect(d, da + 2, db + 2, sfa + 4, sfb + 2, 1u); + // A commit per unit, although the epilogue waits only for the last unit of each chunk: with a commit per + // chunk instead, the issue of a unit takes about 1.4x longer. + commitElect(&accFull[i]); + ++nt; + ++runIdx; + aBase += kTileBytesA; + d += stride; + } + } + + tcFenceBefore(); + __syncthreads(); + if (warp == 5) + { + tcFenceAfter(); + tmemDealloc(tmem, kTmemCols); + } +#endif +} + +void encodeWeightMap(CUtensorMap* map, void const* weights, uint64_t rows, uint32_t boxRows) +{ + cuuint64_t const dims[2] = {64, rows}; + cuuint64_t const strides[1] = {64}; + cuuint32_t const box[2] = {64, boxRows}; + cuuint32_t const elementStrides[2] = {1, 1}; + TLLM_CU_CHECK(tensorrt_llm::common::CUDADriverWrapper::getInstance()->cuTensorMapEncodeTiled(map, + CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(weights), dims, strides, box, elementStrides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_64B, CU_TENSOR_MAP_L2_PROMOTION_L2_128B, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); +} + +} // namespace + +bool isSupported(int32_t hiddenSize, int32_t maxNumCtas) +{ + int const sm = tensorrt_llm::common::getSMVersion(); + if (sm < 100 || sm >= 110 || hiddenSize % kTileRows != 0 || hiddenSize / kTileRows < kMaxUnits + || maxNumCtas < 1 || maxNumCtas > kMaxTokenTiles) + { + return false; + } + int const grid = tensorrt_llm::common::getMultiProcessorCount(); + int const maxUnits = maxNumCtas * (hiddenSize / kTileRows); + return (maxUnits + grid - 1) / grid <= kMaxUnits; +} + +void run(Data const& data, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(isSupported(data.hiddenSize, data.maxNumCtas), "gemm2K128: unsupported shape."); + TLLM_CUDA_CHECK(cudaFuncSetAttribute(gemm2K128Kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes)); + Params p{}; + uint64_t const rows = static_cast(data.numLocalExperts) * data.hiddenSize; + encodeWeightMap(&p.tmW2, data.weights, rows, 2 * kTileRows); + encodeWeightMap(&p.tmW1, data.weights, rows, kTileRows); + p.wSf = static_cast(data.weightsScale); + p.act = static_cast(data.act); + p.actSf = static_cast(data.actScale); + p.scale = data.outputScale; + p.numNonExitingCtas = data.numNonExitingCtas; + p.ctaIdxXyToBatchIdx = data.ctaIdxXyToBatchIdx; + p.ctaIdxXyToMnLimit = data.ctaIdxXyToMnLimit; + p.out = static_cast<__nv_bfloat16*>(data.output); + p.hidden = data.hiddenSize; + p.maxCtas = data.maxNumCtas; + + cudaLaunchConfig_t config{}; + config.gridDim = dim3(tensorrt_llm::common::getMultiProcessorCount()); + config.blockDim = dim3(kThreads); + config.dynamicSmemBytes = kSmemBytes; + config.stream = stream; + cudaLaunchAttribute attributes[2]{}; + attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attributes[0].val.programmaticStreamSerializationAllowed = data.usePdl ? 1 : 0; + attributes[1].id = cudaLaunchAttributePreferredSharedMemoryCarveout; + attributes[1].val.sharedMemCarveout = static_cast(data.sharedMemCarveout); + config.attrs = attributes; + config.numAttrs = data.sharedMemCarveout >= 0 ? 2 : 1; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, gemm2K128Kernel, p)); +} + +} // namespace moe::dev::gemm2K128 diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/Gemm2K128.h b/cpp/tensorrt_llm/kernels/moe/trtllmGen/Gemm2K128.h new file mode 100644 index 000000000000..72d99550e18f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/Gemm2K128.h @@ -0,0 +1,60 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +namespace moe::dev::gemm2K128 +{ + +// FC2 of the NVFP4 MoE at intermediate size 128 and tile N 8, with the inputs and output of the TRTLLM-Gen Gemm2 +// cubins: +// output[8 t + r][f] = bf16(outputScale[e] * sum_k weights[e][f][k] * act[8 t + r][k]), e = ctaIdxXyToBatchIdx[t], +// for the routing's token tiles t < numNonExitingCtas and their valid rows r < ctaIdxXyToMnLimit[t] - 8 t. +struct Data +{ + // [numLocalExperts * hiddenSize][64 B] E2M1, the Gemm2 weight layout (rows shuffled for the 128-row epilogue tile). + void const* weights; + // E4M3 block scales of the weights, 128x4 layout: [numLocalExperts][hiddenSize / 128][2][512 B]. + void const* weightsScale; + // FC1 output: [rows][64 B] E2M1 and its E4M3 block scales in the 8x4 layout ([rows / 8][2][8][4]). + void const* act; + void const* actScale; + float const* outputScale; // [numLocalExperts] + int32_t const* numNonExitingCtas; + int32_t const* ctaIdxXyToBatchIdx; + int32_t const* ctaIdxXyToMnLimit; + void* output; // [rows][hiddenSize] BF16 + int32_t hiddenSize; + int32_t numLocalExperts; + int32_t maxNumCtas; // entries of ctaIdxXyToBatchIdx and ctaIdxXyToMnLimit + bool usePdl; + // The launch's preferred shared-memory carveout in percent (cudaLaunchAttributePreferredSharedMemoryCarveout; + // cudaSharedmemCarveoutMaxShared = 100), or -1: no preference, the driver chooses. + int32_t sharedMemCarveout = -1; +}; + +// Whether run() covers the shape on this device: an SM100-family GPU, hiddenSize a multiple of 128 and at least 1,920, +// at most 64 token tiles, and at most 15 units (token tile x 128 output features) per SM. +bool isSupported(int32_t hiddenSize, int32_t maxNumCtas); + +void run(Data const& data, cudaStream_t stream); + +} // namespace moe::dev::gemm2K128 diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.cu b/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.cu index 4fe30ea94507..cecc50daa877 100644 --- a/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.cu +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.cu @@ -15,6 +15,7 @@ */ #include "DevKernel.h" +#include "Gemm2K128.h" #include "routing/RoutingKernel.h" #include "runner.h" #include "tensorrt_llm/common/config.h" @@ -713,6 +714,10 @@ Runner::Runner(btg::Dtype dtypeAct, btg::Dtype dtypeWeights, bool useDeepSeekFp8 , mGemm2(Gemm2::Runner(dtypeAct, dtypeWeights, btg::Dtype::Bfloat16, useDeepSeekFp8, tileTokensDim, useFineGrained)) , mActType(actType) , mTileTokensDim(tileTokensDim) + , mDtypeAct(dtypeAct) + , mDtypeWeights(dtypeWeights) + , mUseDeepSeekFp8(useDeepSeekFp8) + , mUseFineGrained(useFineGrained) { auto const& gemm1PassingIndices = mPermuteGemm1.getPassingConfigIndices(); auto const& gemm2PassingIndices = mGemm2.getPassingConfigIndices(); @@ -906,13 +911,50 @@ void Runner::run( } // Run gemm2 - mGemm2.run(gemm2_input, gemm2_input_scale, args.gemm2_weights, args.gemm2_weights_scale, args.output2_scales_scalar, - args.gemm2_bias, workspace.gemm2_output, workspace.gemm2_output_scale, totalExpertsPerToken, - args.output_hidden_size.value_or(args.hidden_size), args.intermediate_size, totalLocalExperts, args.num_tokens, - workspace.num_non_exiting_ctas, workspace.total_num_padded_tokens, workspace.cta_idx_xy_to_batch_idx, - workspace.cta_idx_xy_to_mn_limit, workspace.bmm2_workspace, device, stream, config.gemm2Config, - args.valid_hidden_size.value_or(args.hidden_size), - args.valid_intermediate_size.value_or(args.intermediate_size)); + int32_t const gemm2HiddenSize = args.output_hidden_size.value_or(args.hidden_size); + int32_t const gemm2MaxNumCtas + = Routing::getMaxNumCgasInBatchDim(args.num_tokens, totalExpertsPerToken, totalLocalExperts, mTileTokensDim); + // gemm2K128 covers the FC2 of the NVFP4 MoE at intermediate size 128 and tile N 8 (for example GLM-5 decode at MoE + // TP16): the same output rows as the Gemm2 cubin, bitwise (see Gemm2K128.cu). + if (args.prefer_gemm2_k128 && mDtypeAct == btg::Dtype::E2m1 && mDtypeWeights == btg::Dtype::E2m1 + && !mUseDeepSeekFp8 && !mUseFineGrained && mTileTokensDim == 8 && args.num_fused_shared_experts == 0 + && args.intermediate_size == 128 + && args.valid_intermediate_size.value_or(args.intermediate_size) == args.intermediate_size + && args.valid_hidden_size.value_or(args.hidden_size) == gemm2HiddenSize && args.gemm2_bias == nullptr + && args.output2_scales_scalar != nullptr && moe::dev::gemm2K128::isSupported(gemm2HiddenSize, gemm2MaxNumCtas)) + { + moe::dev::gemm2K128::Data gemm2Data{}; + gemm2Data.weights = args.gemm2_weights; + gemm2Data.weightsScale = args.gemm2_weights_scale; + gemm2Data.act = gemm2_input; + gemm2Data.actScale = gemm2_input_scale; + gemm2Data.outputScale = args.output2_scales_scalar; + gemm2Data.numNonExitingCtas = workspace.num_non_exiting_ctas; + gemm2Data.ctaIdxXyToBatchIdx = workspace.cta_idx_xy_to_batch_idx; + gemm2Data.ctaIdxXyToMnLimit = workspace.cta_idx_xy_to_mn_limit; + gemm2Data.output = workspace.gemm2_output; + gemm2Data.hiddenSize = gemm2HiddenSize; + gemm2Data.numLocalExperts = totalLocalExperts; + gemm2Data.maxNumCtas = gemm2MaxNumCtas; + gemm2Data.usePdl = tensorrt_llm::common::getEnvEnablePDL(); + // The maximum shared-memory carveout preferred. Without a preference the driver can run the 141 KB CTAs in the + // 164 KB configuration. The finalize kernel that gemm2K128 lets launch early (for example an all-reduce hop + // with the finalize fused) then shares those SMs and keeps them in that configuration until it ends, so the + // next layer's large-shared-memory CTAs that land there start only after it. A scheduling hint only: the + // outputs are unchanged. + gemm2Data.sharedMemCarveout = static_cast(cudaSharedmemCarveoutMaxShared); + moe::dev::gemm2K128::run(gemm2Data, stream); + } + else + { + mGemm2.run(gemm2_input, gemm2_input_scale, args.gemm2_weights, args.gemm2_weights_scale, + args.output2_scales_scalar, args.gemm2_bias, workspace.gemm2_output, workspace.gemm2_output_scale, + totalExpertsPerToken, gemm2HiddenSize, args.intermediate_size, totalLocalExperts, args.num_tokens, + workspace.num_non_exiting_ctas, workspace.total_num_padded_tokens, workspace.cta_idx_xy_to_batch_idx, + workspace.cta_idx_xy_to_mn_limit, workspace.bmm2_workspace, device, stream, config.gemm2Config, + args.valid_hidden_size.value_or(args.hidden_size), + args.valid_intermediate_size.value_or(args.intermediate_size)); + } // Run finalize if (args.do_finalize) diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.h b/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.h index cbe79c27e961..47e5a5b98e60 100644 --- a/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.h +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.h @@ -366,6 +366,10 @@ struct MoERunnerArgs // finalize bool do_finalize{true}; + + // Run gemm2 on moe::dev::gemm2K128 instead of the Gemm2 cubin where it covers the shape (see Runner::run); the + // output rows are the cubin's, bitwise. + bool prefer_gemm2_k128{false}; }; struct MoEWorkspace @@ -455,6 +459,11 @@ class Runner // Kept so setOpsData can tell the activation launcher how much per-expert // tile padding the permuted row space carries. int32_t mTileTokensDim; + // Kept to tell whether gemm2K128 covers this runner's Gemm2 (see run()). + batchedGemm::trtllm::gen::Dtype mDtypeAct; + batchedGemm::trtllm::gen::Dtype mDtypeWeights; + bool mUseDeepSeekFp8; + bool mUseFineGrained; // This will be the cartesian product of the passing configs for gemm1 and gemm2 // This allows us to autotune the MoE as one operation instead of tuning gemm1 and gemm2 separately diff --git a/cpp/tensorrt_llm/thop/moe/trtllmGen/fp4BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/moe/trtllmGen/fp4BlockScaleMoe.cpp index 4b0859f02175..5333d5daa8dd 100644 --- a/cpp/tensorrt_llm/thop/moe/trtllmGen/fp4BlockScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/moe/trtllmGen/fp4BlockScaleMoe.cpp @@ -50,7 +50,7 @@ std::vector run_fp4_block_scale_moe_runner(torch::optional const hidden_states_ready_event = std::nullopt, torch::optional const& router_hidden_states = torch::nullopt, torch::optional const& router_weight = torch::nullopt, - torch::optional const& router_workspace = torch::nullopt) + torch::optional const& router_workspace = torch::nullopt, bool const prefer_gemm2_k128 = false) { TORCH_CHECK(dtype == btg::Dtype::E4m3 || dtype == btg::Dtype::E2m1, "dtype can only be e4m3 or e2m1."); TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "Only SM100f is supported by FP4 block scale MOE"); @@ -528,6 +528,7 @@ std::vector run_fp4_block_scale_moe_runner(torch::optional(); args.output2_scales_scalar = output2_scales_scalar.data_ptr(); args.do_finalize = do_finalize; + args.prefer_gemm2_k128 = prefer_gemm2_k128; auto const workspace_sizes = moe_runner.getWorkspaceSizeInBytes(args, moeConfigIndex); @@ -635,7 +636,7 @@ class FP4BlockScaleMoeRunner : public torch::CustomClassHolder std::optional const hidden_states_ready_event = std::nullopt, torch::optional const& router_hidden_states = torch::nullopt, torch::optional const& router_weight = torch::nullopt, - torch::optional const& router_workspace = torch::nullopt) + torch::optional const& router_workspace = torch::nullopt, bool const prefer_gemm2_k128 = false) { int64_t tileN = moeConfigIndex[0]; int64_t config = moeConfigIndex[1]; @@ -686,7 +687,7 @@ class FP4BlockScaleMoeRunner : public torch::CustomClassHolder num_experts, top_k, n_group, topk_group, intermediate_size, local_expert_offset, local_num_experts, routed_scaling_factor, tileN, routing_method_type, do_finalize, mDtypeElt, *runner, configToRun, topk_weights, topk_ids, output, hidden_states_ready_event, router_hidden_states, router_weight, - router_workspace); + router_workspace, prefer_gemm2_k128); } private: diff --git a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py index 6e43eacb9078..b973aeb0bcfc 100644 --- a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py @@ -461,7 +461,8 @@ def __init__(self, do_finalize: bool, act_type: int, tune_max_num_tokens: int = 8192, - use_dp: bool = False): + use_dp: bool = False, + prefer_gemm2_k128: bool = False): self.num_experts = num_experts self.top_k = top_k @@ -474,6 +475,7 @@ def __init__(self, self.routing_method_type = routing_method_type self.do_finalize = do_finalize self.act_type = act_type + self.prefer_gemm2_k128 = prefer_gemm2_k128 self.tuning_config = FP4BlockScaleMoERunner.get_tuning_config( self.num_experts // self.local_num_experts, @@ -483,8 +485,10 @@ def __init__(self, # The unique_id is used by the autotuner to get the cache key, so we hash on members # that influence tactic validity here. e.g. we are tuning FC1 and FC2 so the routing type does not matter def unique_id(self): - return (self.top_k, self.intermediate_size, self.local_num_experts, - self.act_type) + key = (self.top_k, self.intermediate_size, self.local_num_experts, + self.act_type) + # FC2 on gemm2K128 times the tile-8 tactics differently. + return key + ("gemm2_k128", ) if self.prefer_gemm2_k128 else key def get_runner(self): instance_key = (self.act_type, ) @@ -522,7 +526,7 @@ def forward( self.routed_scaling_factor, self.routing_method_type, self.do_finalize, tactic, args.topk_weights, args.topk_ids, output, hidden_states_ready_event, router_hidden_states, router_weight, - router_workspace) + router_workspace, self.prefer_gemm2_k128) def get_valid_tactics(self, inputs: List[torch.Tensor], profile: OptimizationProfile, @@ -690,7 +694,8 @@ def fp4_block_scale_moe_runner(routing_logits: Optional[torch.Tensor], router_hidden_states: Optional[ torch.Tensor] = None, router_weight: Optional[torch.Tensor] = None, - router_workspace: Optional[torch.Tensor] = None + router_workspace: Optional[torch.Tensor] = None, + prefer_gemm2_k128: bool = False ) -> List[torch.Tensor]: """hidden_states_ready_event: raw cudaEvent_t (torch.cuda.Event.cuda_event) recorded on another stream once hidden_states and hidden_states_scale are @@ -703,7 +708,12 @@ def fp4_block_scale_moe_runner(routing_logits: Optional[torch.Tensor], there and routing_logits only gives their shape (it is not read or written). router_workspace: int32 [FUSED_ROUTER_WORKSPACE_WORDS], see make_fused_router_workspace; the kernel uses and restores it. Profiling - runs of the autotuner route from routing_logits as usual.""" + runs of the autotuner route from routing_logits as usual. + + prefer_gemm2_k128: FC2 runs on the persistent K = 128 kernel where it + covers the shape (NVFP4, tile N 8, intermediate size 128, no bias, enough + SMs), with the output rows of the Gemm2 cubin, bitwise; elsewhere the cubin + runs.""" tuner = AutoTuner.get() kernel_runner = FP4BlockScaleMoERunner( @@ -720,6 +730,7 @@ def fp4_block_scale_moe_runner(routing_logits: Optional[torch.Tensor], act_type, tune_max_num_tokens=tune_max_num_tokens, use_dp=use_dp, + prefer_gemm2_k128=prefer_gemm2_k128, ) # Prepare dummy topk tensors and hook for AutoTuner profiling @@ -873,7 +884,8 @@ def _(routing_logits, hidden_states_ready_event: Optional[int] = None, router_hidden_states: Optional[torch.Tensor] = None, router_weight: Optional[torch.Tensor] = None, - router_workspace: Optional[torch.Tensor] = None) -> List[torch.Tensor]: + router_workspace: Optional[torch.Tensor] = None, + prefer_gemm2_k128: bool = False) -> List[torch.Tensor]: if do_finalize: num_tokens = hidden_states.shape[0] hidden_size = hidden_states.shape[1] * 2 From 84f8e7e55ebb909c8b47e6a8a8de82facfd66d77 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:05:00 -0700 Subject: [PATCH 17/53] Add the GLM-5.2 tp16 routed-experts entry ops/moe/routed_experts wraps the fp4_block_scale_moe_runner call the target makes in its MoE layers, every argument explicit: DeepSeek-V3 routing without groups over all 256 experts at MoE TP16 (intermediate 128), with the router GEMV inside the routing kernel at up to 8 rows, the NVFP4 input joined right before FC1 and FC2 on the K = 128 kernel, unfinalized for the post-MoE hop. Cells: the 8-row verify step, one row, and 32 rows routed from the router GEMM's logits. The reference is a plain-torch NVFP4 dequant MoE; is_valid refuses strided operands, an unrecorded event and a router workspace not in its between-launch state. The sm_107-only test drives the cells against the reference, with the input quantized on a side stream behind a sleep, and proves the fused router bitwise against the router GEMM at 1-8 rows, the K = 128 FC2 bitwise against the cubin, the three-kernel launch and CUDA-graph replay. The catalog fp4_quantize entry gains the target's 8-row cell at hidden 6144, its routed-expert input. Signed-off-by: Dongfeng Yu --- .../catalog/quantization/fp4_quantize.py | 12 +- .../ops/moe/routed_experts.py | 414 +++++++++++ .../moe/test_modeling_v2_routed_experts.py | 650 ++++++++++++++++++ 3 files changed, 1072 insertions(+), 4 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/routed_experts.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_routed_experts.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py index b0c0e864d886..2e23c1c3c997 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py @@ -51,10 +51,10 @@ class Fp4Quantize(OpWrapper): # nvfp4_gemm in the dense and shared-expert MLPs, the linear one feeding the # routed-expert path, at its hidden size and its intermediate. qwen3.8 # (qwen4_exp) calls it only for its routed experts: linear scales at hidden - # 2560. A linear buffer has no row padding, so at a row count that is not a - # multiple of 128 it is shorter than the swizzled one; at a multiple of 128 - # the two are the same length and only the cells' byte order tells them - # apart. + # 2560; glm-5.2 likewise, at hidden 6144. A linear buffer has no row + # padding, so at a row count that is not a multiple of 128 it is shorter + # than the swizzled one; at a multiple of 128 the two are the same length + # and only the cells' byte order tells them apart. CELLS: tuple[Cell, ...] = ( Cell( why="decode, hidden 7168, swizzled scales -- the dense MLP's input", @@ -88,6 +88,10 @@ class Fp4Quantize(OpWrapper): why="qwen3.8 prefill at max_num_tokens 8192, hidden 2560, linear", spec=dict(shape=(8192, 2560), dtype=torch.bfloat16, vec=16, swizzled=False), ), + Cell( + why="glm-5.2 verify step: 8 rows at hidden 6144, linear -- its routed-expert input", + spec=dict(shape=(8, 6144), dtype=torch.bfloat16, vec=16, swizzled=False), + ), ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/routed_experts.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/routed_experts.py new file mode 100644 index 000000000000..5ff359826b89 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/routed_experts.py @@ -0,0 +1,414 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GLM-5.2's routed experts at MoE TP16: DeepSeek-V3 routing without groups +(sigmoid + bias, top-8 of 256, scaled sum normalization), grouped NVFP4 FC1 + +SwiGLU + NVFP4 requantization + grouped NVFP4 FC2, in one trtllm-gen call. + +On the decode step (up to 8 rows) the call runs the router GEMV inside the +routing kernel, lets the NVFP4 input still be in flight on another stream +until FC1 starts, and runs FC2 on the K = 128 kernel; it stops before the +finalize step, which the post-MoE all-reduce hop performs.""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( # noqa: E501 + dequantize_nvfp4, + quantize_nvfp4, +) +from tensorrt_llm._torch.custom_ops.trtllm_gen_custom_ops import ( + FUSED_ROUTER_WORKSPACE_WORDS, + make_fused_router_workspace, +) + +__all__ = [ + "ACT_TYPE_SWIGLU", + "FUSED_ROUTER_MAX_TOKENS", + "FUSED_ROUTER_WORKSPACE_WORDS", + "ROUTING_METHOD_DEEPSEEK_V3", + "make_fused_router_workspace", + "routed_experts", +] + +#: `routing_method_type` of DeepSeek-V3 routing; with `n_group` 1 the op runs +#: it as sigmoid + bias, top-k, and the sum-normalized, scaled sigmoid weights. +ROUTING_METHOD_DEEPSEEK_V3 = 2 + +#: `act_type` for SwiGlu. +ACT_TYPE_SWIGLU = 0 + +#: The largest row count the fused router + routing kernel takes. +FUSED_ROUTER_MAX_TOKENS = 8 + +#: One expert stack's `(codes [E, N, K] uint8, block scales [E, N, K/16] +#: float8_e4m3fn)`, unshuffled and unswizzled. +Operand = Tuple[torch.Tensor, torch.Tensor] + +_IDLE_BITS = -(2**31) # -0.0f, the idle value of a fused-router exchange buffer + + +def _dequantize_activation(hidden_states: torch.Tensor, hidden_states_scale: torch.Tensor): + """`[T, H]` fp32 value of an NVFP4 activation in the op's input layout: + packed e2m1 codes `[T, H/2]` (element 2i in the low nibble) and linear e4m3 + block scales, one per 16 elements of a row. The global scale stays in.""" + rows = hidden_states.shape[0] + codes = torch.stack([hidden_states & 0x0F, hidden_states >> 4], dim=-1).reshape(rows, -1) + scales = hidden_states_scale.view(torch.float8_e4m3fn).reshape(rows, -1) + return dequantize_nvfp4(codes, scales) + + +class GlmRoutedExperts(OpWrapper): + """One routed-expert MoE layer of GLM-5.2 at MoE TP16, over NVFP4 input. + + Routing happens inside the op, from one of two sources: + + * `routing_logits` `[T, 256]` -- the router GEMM's fp32 logits; + * `router_hidden_states` `[T, 6144]` bf16 with `router_weight` + `[256, 6144]` bf16 and `router_workspace` (`make_fused_router_workspace`): + the routing kernel runs the router GEMV itself, bit for bit the logits of + `dsv3_router_gemm_op`, at 1-8 rows. `routing_logits` then only gives the + shape and is not read. + + `hidden_states_ready_event` is the raw `cudaEvent_t` of a `torch.cuda.Event` + recorded on the stream that writes `hidden_states` / `hidden_states_scale` + (`event.cuda_event`, after its first record): routing launches without + waiting for it, FC1 waits. `prefer_gemm2_k128` runs FC2 on the persistent + K = 128 kernel where it covers the call -- intermediate 128, tile N 8 and + enough SMs, which an 8-row call meets on Rubin's SM count -- with the Gemm2 + cubin's rows, bitwise. + + `do_finalize=False` returns `[fc2 rows [P, H] bf16, expert weights [T, 8] + bf16, expanded index -> permuted row [T, 8] int32]`: token `t`'s output is + `sum_k weights[t, k] * rows[map[t, k]]`. `do_finalize=True` returns + `[output [T, H] bf16]`. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # GLM-5.2's routed geometry at MoE TP16, ep 1: every rank holds all 256 + # experts with a 2048 / 16 = 128-wide intermediate slice. + _GLM_TP16 = dict( + hidden=6144, + intermediate=128, + num_experts=256, + top_k=8, + routed_scaling_factor=2.5, + local_expert_offset=0, + local_num_experts=256, + ) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="the c1 verify step: 8 rows routed by the fused router GEMV, the input joined " + "right before FC1, FC2 at K = 128, unfinalized for the post-MoE hop", + spec=dict( + tokens=8, fused_router=True, late_join=True, gemm2_k128=True, **_GLM_TP16 + ), + ), + Cell( + why="one row, a step without MTP: the fused router's single-token instance and " + "FC2 at K = 128 over 8 token tiles", + spec=dict( + tokens=1, fused_router=True, late_join=True, gemm2_k128=True, **_GLM_TP16 + ), + ), + Cell( + why="32 rows (4 requests x 8), past the fused router: routed from the router " + "GEMM's logits, FC2 back on the Gemm2 cubin", + spec=dict( + tokens=32, fused_router=False, late_join=False, gemm2_k128=True, **_GLM_TP16 + ), + ), + ) + + note = """ + Certified unfinalized (`do_finalize=False`) only; the finalize step is the + post-MoE hop's. No expert-parallel window (`local_expert_offset` 0, all 256 + experts local), no `output` buffer, `use_dp` False, `topk_ids` / + `topk_weights` None. `gemm1_bias`, `gemm1_alpha`, `gemm1_beta`, + `gemm1_clamp_limit` and `gemm2_bias` are None, `act_type` SwiGlu. + + `routing_bias` is GLM-5.2's `e_score_correction_bias` in bf16, the dtype + the TRTLLM MoE backend loads it in; `n_group` / `topk_group` are 1. + + The tactic is the autotuner's: the cells run cold (the op's default tile, + N 8 at these row counts). Under autotuning, a call routed by the fused + router profiles from `routing_logits`, so the caller fills that buffer + once (random logits) rather than leaving it uninitialized. + + A router workspace is per call site and its launches must run in stream + order; the kernel leaves the exchange buffer of the next launch idle and + advances the epoch word. `tune_max_num_tokens` is an autotuner bucket hint. + + The caller owns the weight layout -- FC1 rows `[up ; gate]` interleaved and + 32-row block shuffled, FC2 rows block shuffled, both scale stacks 128x4 + swizzled -- as the catalog `fp4_block_scale_moe_runner` entry describes. + """ + + def __call__( + self, + routing_logits: torch.Tensor, + routing_bias: torch.Tensor, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm1_bias: Optional[torch.Tensor], + gemm1_alpha: Optional[torch.Tensor], + gemm1_beta: Optional[torch.Tensor], + gemm1_clamp_limit: Optional[torch.Tensor], + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + gemm2_bias: Optional[torch.Tensor], + output1_scale_scalar: torch.Tensor, + output1_scale_gate_scalar: torch.Tensor, + output2_scale_scalar: torch.Tensor, + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: float, + routing_method_type: int, + do_finalize: bool, + act_type: int, + topk_weights: Optional[torch.Tensor], + topk_ids: Optional[torch.Tensor], + output: Optional[torch.Tensor], + tune_max_num_tokens: int, + use_dp: bool, + hidden_states_ready_event: Optional[int], + router_hidden_states: Optional[torch.Tensor], + router_weight: Optional[torch.Tensor], + router_workspace: Optional[torch.Tensor], + prefer_gemm2_k128: bool, + ) -> list[torch.Tensor]: + return torch.ops.trtllm.fp4_block_scale_moe_runner( + routing_logits, + routing_bias, + hidden_states, + hidden_states_scale, + gemm1_weights, + gemm1_weights_scale, + gemm1_bias, + gemm1_alpha, + gemm1_beta, + gemm1_clamp_limit, + gemm2_weights, + gemm2_weights_scale, + gemm2_bias, + output1_scale_scalar, + output1_scale_gate_scalar, + output2_scale_scalar, + num_experts, + top_k, + n_group, + topk_group, + intermediate_size, + local_expert_offset, + local_num_experts, + routed_scaling_factor, + routing_method_type, + do_finalize, + act_type, + topk_weights=topk_weights, + topk_ids=topk_ids, + output=output, + tune_max_num_tokens=tune_max_num_tokens, + use_dp=use_dp, + hidden_states_ready_event=hidden_states_ready_event, + router_hidden_states=router_hidden_states, + router_weight=router_weight, + router_workspace=router_workspace, + prefer_gemm2_k128=prefer_gemm2_k128, + ) + + def reference( + self, + routing_logits: Optional[torch.Tensor], + routing_bias: torch.Tensor, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + up: Operand, + gate: Operand, + down: Operand, + output1_scale_scalar: torch.Tensor, + output1_scale_gate_scalar: torch.Tensor, + output2_scale_scalar: torch.Tensor, + top_k: int, + local_expert_offset: int, + routed_scaling_factor: float, + router_hidden_states: Optional[torch.Tensor] = None, + router_weight: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """The layer, in plain torch over dequantized weights, fp32 throughout. + + Returns `(output [T, H] bf16, topk ids [T, top_k], topk weights [T, + top_k] bf16)`: the finalized layer, rounded once, and the routing it + used. A `do_finalize=False` call is compared after summing its rows with + its own weights. + + The documented exception to mirroring: `__call__` takes the prepared + weight operands -- shuffled rows, swizzled scales -- and a reference + given those would re-derive the layout rather than the computation, so + this takes the per-expert `(codes, block scales)` pairs `up`, `gate` + and `down` the layout was built from. Everything else keeps the name it + has in `__call__`. + + Routing: the logits are `routing_logits`, or the router GEMV + `router_hidden_states @ router_weight.T` when those are given (the + fused router). Scores are `sigmoid(logits)` -- in the routing kernel's + `0.5 tanh(x / 2) + 0.5` form --, experts are chosen by score + bias, + and an expert's weight is its score over the sum of the chosen scores, + times `routed_scaling_factor`, rounded to bf16 as the kernel stores it. + + Experts: the op's own scalars carry the global scales. FC1's gate half + is dequantized by `output1_scale_gate_scalar` (`a1`), its linear half + by the same `a1`; the SwiGLU output is requantized to NVFP4 under the + FC2-input global scale `g2 = output1_scale_scalar / a1`, which the FC1 + epilogue folds in, and FC2's accumulator is dequantized by + `output2_scale_scalar`. Skipping that requantization is a reference + some 20 ulp RMS away from a correct kernel. + """ + x = _dequantize_activation(hidden_states, hidden_states_scale) + num_tokens, hidden = x.shape + with true_fp32_matmul(): + if router_hidden_states is not None: + logits = router_hidden_states.float() @ router_weight.float().t() + else: + logits = routing_logits.float() + scores = 0.5 * torch.tanh(0.5 * logits) + 0.5 + chosen = torch.topk(scores + routing_bias.float(), top_k, dim=-1).indices + picked = scores.gather(1, chosen) + weights = (picked * routed_scaling_factor / picked.sum(-1, keepdim=True)).to( + torch.bfloat16 + ) + + up_c, up_s = up + gt_c, gt_s = gate + dn_c, dn_s = down + out = torch.zeros(num_tokens, hidden, dtype=torch.float32, device=x.device) + for local_e in range(up_c.shape[0]): + tok, slot = (chosen == local_expert_offset + local_e).nonzero(as_tuple=True) + if tok.numel() == 0: + continue + a1 = output1_scale_gate_scalar[local_e].float() + g2 = output1_scale_scalar[local_e].float() / a1 + a2 = output2_scale_scalar[local_e].float() + xe = x[tok] + u = (xe @ dequantize_nvfp4(up_c[local_e], up_s[local_e]).t()) * a1 + g = (xe @ dequantize_nvfp4(gt_c[local_e], gt_s[local_e]).t()) * a1 + act = quantize_nvfp4(u * g * torch.sigmoid(g), float(g2)) + y = (act @ dequantize_nvfp4(dn_c[local_e], dn_s[local_e]).t()) * a2 + out.index_add_(0, tok, y * weights[tok, slot].float().unsqueeze(1)) + return out.to(torch.bfloat16), chosen.to(torch.int32), weights + + def is_valid( + self, + routing_logits: torch.Tensor, + routing_bias: torch.Tensor, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm1_bias: Optional[torch.Tensor], + gemm1_alpha: Optional[torch.Tensor], + gemm1_beta: Optional[torch.Tensor], + gemm1_clamp_limit: Optional[torch.Tensor], + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + gemm2_bias: Optional[torch.Tensor], + output1_scale_scalar: torch.Tensor, + output1_scale_gate_scalar: torch.Tensor, + output2_scale_scalar: torch.Tensor, + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: float, + routing_method_type: int, + do_finalize: bool, + act_type: int, + topk_weights: Optional[torch.Tensor], + topk_ids: Optional[torch.Tensor], + output: Optional[torch.Tensor], + tune_max_num_tokens: int, + use_dp: bool, + hidden_states_ready_event: Optional[int], + router_hidden_states: Optional[torch.Tensor], + router_weight: Optional[torch.Tensor], + router_workspace: Optional[torch.Tensor], + prefer_gemm2_k128: bool, + ) -> None: + """The ways this op takes an input and answers it wrongly. + + Shapes, dtypes, the routing configuration the fused router covers and + the workspace size are rejected loudly by the op and are not restated. + """ + # The kernels take raw data pointers and assume dense row-major + # tensors; a strided view is read as if dense. + for name, tensor in ( + ("routing_logits", routing_logits), + ("routing_bias", routing_bias), + ("hidden_states", hidden_states), + ("hidden_states_scale", hidden_states_scale), + ("gemm1_weights", gemm1_weights), + ("gemm1_weights_scale", gemm1_weights_scale), + ("gemm2_weights", gemm2_weights), + ("gemm2_weights_scale", gemm2_weights_scale), + ("output1_scale_scalar", output1_scale_scalar), + ("output1_scale_gate_scalar", output1_scale_gate_scalar), + ("output2_scale_scalar", output2_scale_scalar), + ("router_hidden_states", router_hidden_states), + ("router_weight", router_weight), + ("router_workspace", router_workspace), + ("output", output), + ): + if tensor is not None: + assert tensor.is_contiguous(), ( + f"{name} must be contiguous; a strided view is read as if dense " + "and silently produces wrong results" + ) + # The prepared block-scale layout is misread at a hidden size that is + # only a multiple of 128 (see the catalog fp4_block_scale_moe_runner). + assert gemm1_weights.shape[-1] * 2 % 256 == 0, ( + "hidden size (gemm1_weights.shape[-1] * 2) must be a multiple of 256; " + "the prepared block-scale layout is silently misread otherwise" + ) + # A torch.cuda.Event has no cudaEvent_t before its first record, and a + # null handle reads as "no event": FC1 would not wait for the input. + assert hidden_states_ready_event is None or hidden_states_ready_event != 0, ( + "hidden_states_ready_event is a null cudaEvent_t: take event.cuda_event after " + "the event is recorded, or FC1 reads hidden_states without waiting" + ) + # The routing CTA takes any value but -0.0 in the exchange buffer of + # this launch as a published slice sum. Unless that buffer is idle, it + # routes on stale sums -- a zero-filled workspace routes on zero logits. + if router_hidden_states is not None and router_workspace is not None: + words = router_workspace.view(torch.int32) + buffer_words = (FUSED_ROUTER_WORKSPACE_WORDS - 4) // 2 + epoch = int(words[2 * buffer_words].item()) + current = words[(epoch & 1) * buffer_words : ((epoch & 1) + 1) * buffer_words] + assert bool((current == _IDLE_BITS).all()), ( + "router_workspace is not in its between-launch state (the exchange buffer " + "of the next launch must hold -0.0 everywhere); build it with " + "make_fused_router_workspace and let only this op write it" + ) + + +routed_experts = GlmRoutedExperts() diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_routed_experts.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_routed_experts.py new file mode 100644 index 000000000000..059591716d8f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_routed_experts.py @@ -0,0 +1,650 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM-5.2 tp16 routed_experts entry.""" + +import zlib + +import pytest +import torch +from torch.profiler import ProfilerActivity, profile + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.moe.routed_experts import ( # noqa: E501 + ACT_TYPE_SWIGLU, + FUSED_ROUTER_WORKSPACE_WORDS, + ROUTING_METHOD_DEEPSEEK_V3, + make_fused_router_workspace, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.moe.routed_experts import ( # noqa: E501 + routed_experts as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402,F401 — needs the path declared above + assert_within_ulp, + certifying, + validating, +) + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +DEV = "cuda" +HIDDEN, INTER, EXPERTS, TOP_K = 6144, 128, 256, 8 +ROUTED_SCALE = 2.5 +SV = 16 # NVFP4 block size: one e4m3 scale per 16 elements along K +E2M1_MAX, E4M3_MAX = 6.0, 448.0 +# Expert-selection scores (sigmoid + bias) of the 8th and 9th expert closer than +# this leave the reference's choice to the sigmoid's last bit; a case compared +# against the reference is drawn clear of it. +SELECTION_MARGIN = 1e-6 +# Long enough (tens of ms) that a kernel not waiting for the side stream reads +# the stale input long before the side stream writes the real one. +SLEEP_CYCLES = 1 << 27 +_BUFFER_WORDS = (FUSED_ROUTER_WORKSPACE_WORDS - 4) // 2 +_EPOCH_WORD = 2 * _BUFFER_WORDS + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """16 ulp per element against the row's own scale, 2 ulp relative RMS. + + The band of the catalog `fp4_block_scale_moe_runner` entry, which compares + the same kernels against the same reference over bit-identical NVFP4 + operands: they differ only in accumulation order and in whether a marginal + FC1 value lands on the same e2m1 code at the FC2-input requantization. Its + measured worst over every geometry it covers is 9.97 ulp element-wise and + 0.88 ulp RMS, 4.68 / 0.70 for DeepSeek-R1's 256-expert stack; a reference + that skips the requantization lands 23 ulp RMS out. GLM's FC2 sums over 128 + terms where R1's sums over 2048, so accumulation order moves it less. + Routing adds one source those comparisons did not have: a routing weight + the kernel rounds to bf16 one ulp away from the reference's moves one + slot's share of a row by 2^-8, under 1 ulp of the row. A different expert + set is hundreds of ulp out. + """ + assert_within_ulp(out, ref, element_ulp=16.0, rms_ulp=2.0) + + +# ── operands ────────────────────────────────────────────────────────────── + + +def _rand_nvfp4(e: int, n: int, k: int, gen: torch.Generator): + """Random NVFP4 expert stack: `(packed [E, N, K/2] uint8, codes [E, N, K] + uint8, scales [E, N, K/16] float8_e4m3fn)`. Codes and power-of-two scales + are drawn directly, so every weight's fp32 value is exact.""" + codes = torch.randint(0, 16, (e, n, k), dtype=torch.uint8, device=DEV, generator=gen) + exps = torch.randint(-2, 1, (e, n, k // SV), device=DEV, generator=gen, dtype=torch.int32) + scales = torch.exp2(exps.float()).to(torch.float8_e4m3fn) + packed = (codes[..., 0::2] | (codes[..., 1::2] << 4)).contiguous() + return packed, codes, scales + + +def _blk32_perm(m: int) -> torch.Tensor: + """Gather index of the 32-row block shuffle: within each block of 32 rows, + source row `4u + v` lands at destination row `8v + u`.""" + j = torch.arange(32) + idx = torch.empty(32, dtype=torch.long) + idx[(j % 4) * 8 + j // 4] = j + return (idx.repeat(m // 32) + torch.arange(m // 32).repeat_interleave(32) * 32).to(DEV) + + +def _interleave_perm(m: int) -> torch.Tensor: + """Gather index interleaving `[up | gate]` rows into up0, gate0, up1, ...""" + p = torch.empty(m, dtype=torch.long) + p[0::2] = torch.arange(0, m // 2) + p[1::2] = torch.arange(m // 2, m) + return p.to(DEV) + + +def _swizzle(s: torch.Tensor) -> torch.Tensor: + """Per-expert 128x4 block-scale swizzle of `[E, M, C]` uint8 scales.""" + e, m, c = s.shape + return s.reshape(e, m // 128, 4, 32, c // 4, 4).permute(0, 1, 4, 3, 2, 5).reshape(e, m, c) + + +def _prepare(packed: torch.Tensor, scales: torch.Tensor, perm: torch.Tensor): + """Shuffle the rows of a stack and its scales, then swizzle the scales.""" + w = torch.index_select(packed, 1, perm).contiguous() + s = torch.index_select(scales.view(torch.uint8), 1, perm) + return w, _swizzle(s).contiguous().view(torch.float8_e4m3fn) + + +_STACK = None + + +def _stack(): + """GLM-5.2's expert stack at MoE TP16 -- 256 experts, H 6144, I 128 -- + built once: the kernel operands and the reference's `(codes, scales)`.""" + global _STACK + if _STACK is None: + gen = torch.Generator(device=DEV).manual_seed(5216) + up_p, up_c, up_s = _rand_nvfp4(EXPERTS, INTER, HIDDEN, gen) + gt_p, gt_c, gt_s = _rand_nvfp4(EXPERTS, INTER, HIDDEN, gen) + dn_p, dn_c, dn_s = _rand_nvfp4(EXPERTS, HIDDEN, INTER, gen) + fc1_perm = _interleave_perm(2 * INTER)[_blk32_perm(2 * INTER)] + w1, s1 = _prepare( + torch.cat([up_p, gt_p], dim=1), torch.cat([up_s, gt_s], dim=1), fc1_perm + ) + w2, s2 = _prepare(dn_p, dn_s, _blk32_perm(HIDDEN)) + _STACK = dict( + gemm1_weights=w1, + gemm1_weights_scale=s1, + gemm2_weights=w2, + gemm2_weights_scale=s2, + up=(up_c, up_s), + gate=(gt_c, gt_s), + down=(dn_c, dn_s), + ) + return _STACK + + +def _router(gen: torch.Generator): + """Router weight and bias. The weight sits on a 2^-8 grid within +-1/32 so + that, against a router input on a quarter grid, every partial sum of a + router GEMV is exact in fp32: any summation order -- the kernel's, torch's + -- gives the same logits.""" + weight = ( + torch.randint(-8, 9, (EXPERTS, HIDDEN), device=DEV, generator=gen).float() * 2.0**-8 + ).to(torch.bfloat16) + bias = (torch.randn(EXPERTS, device=DEV, generator=gen) * 0.05).to(torch.bfloat16) + return weight, bias + + +def _router_input(tokens: int, gen: torch.Generator) -> torch.Tensor: + """bf16 router input on a quarter grid in [-1, 1] (see `_router`).""" + ints = torch.randint(-4, 5, (tokens, HIDDEN), device=DEV, generator=gen) + return (ints.float() * 0.25).to(torch.bfloat16) + + +def _input_global_scale(x: torch.Tensor) -> torch.Tensor: + """The FC1-input global scale a checkpoint's `input_scale` encodes.""" + return (E4M3_MAX * E2M1_MAX / x.float().abs().amax()).reshape(1).to(torch.float32) + + +def _quantize(x: torch.Tensor, g1: torch.Tensor): + """The experts' NVFP4 input, as the target quantizes it: linear scales.""" + xq, xsf = torch.ops.trtllm.fp4_quantize(x, g1, SV, False, False) + return xq, xsf.view(torch.float8_e4m3fn) + + +def _scalars(xq, xsf, chosen, g1: torch.Tensor): + """The three per-expert scalars for weights with a global scale of 1, and + `g2` calibrated on the true FC1 activation of the chosen experts.""" + stack = _stack() + x = _dequant_act(xq, xsf) / g1 + amax = torch.zeros((), device=DEV) + with true_fp32_matmul(): + for e in chosen.unique().tolist(): + tok = (chosen == e).any(dim=1).nonzero(as_tuple=True)[0] + up = x[tok] @ _dequant(*_slice(stack["up"], e)).t() + gate = x[tok] @ _dequant(*_slice(stack["gate"], e)).t() + amax = torch.maximum(amax, (up * gate * torch.sigmoid(gate)).abs().max()) + g2 = E4M3_MAX * E2M1_MAX / amax + a1 = torch.full((EXPERTS,), 1.0, device=DEV) / g1 + return a1 * g2, a1, torch.full((EXPERTS,), 1.0, device=DEV) / g2 + + +_E2M1 = None + + +def _dequant(codes: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + global _E2M1 + if _E2M1 is None: + _E2M1 = torch.tensor( + [0, 0.5, 1, 1.5, 2, 3, 4, 6, -0.0, -0.5, -1, -1.5, -2, -3, -4, -6], + dtype=torch.float32, + device=DEV, + ) + return _E2M1[codes.long()] * scales.float().repeat_interleave(SV, dim=1) + + +def _slice(operand, e: int): + codes, scales = operand + return codes[e], scales[e] + + +def _dequant_act(xq: torch.Tensor, xsf: torch.Tensor) -> torch.Tensor: + codes = torch.stack([xq & 0x0F, xq >> 4], dim=-1).reshape(xq.shape[0], -1) + return _dequant(codes, xsf.view(torch.float8_e4m3fn).reshape(xq.shape[0], -1)) + + +def _selection(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor): + """`(logits, top-8 by score + bias, smallest 8th-to-9th score gap)`.""" + with true_fp32_matmul(): + logits = x.float() @ weight.float().t() + sel = 0.5 * torch.tanh(0.5 * logits) + 0.5 + bias.float() + top = torch.topk(sel, TOP_K + 1, dim=-1) + margin = (top.values[:, TOP_K - 1] - top.values[:, TOP_K]).min().item() + return logits, top.indices[:, :TOP_K], margin + + +def _case(tokens: int, seed: int, clear_margin: bool = False) -> dict: + """One call's inputs. `clear_margin` redraws (deterministically, from the + next seeds) a router input whose expert choice sits within + `SELECTION_MARGIN` of a tie, which only a comparison against the + reference needs.""" + for attempt in range(16): + gen = torch.Generator(device=DEV).manual_seed(seed + attempt) + weight, bias = _router(gen) + x = _router_input(tokens, gen) + logits, chosen, margin = _selection(x, weight, bias) + if not clear_margin or margin > SELECTION_MARGIN: + break + else: + raise AssertionError(f"no draw from seed {seed} clears the selection margin") + g1 = _input_global_scale(x) + xq, xsf = _quantize(x, g1) + s1, sg, s2 = _scalars(xq, xsf, chosen, g1) + return dict( + x=x, + g1=g1, + xq=xq, + xsf=xsf, + logits=logits, + router_weight=weight, + routing_bias=bias, + output1_scale_scalar=s1, + output1_scale_gate_scalar=sg, + output2_scale_scalar=s2, + ) + + +# ── the call ────────────────────────────────────────────────────────────── + + +def _call( + case: dict, + *, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + fused_router: bool, + workspace=None, + event=None, + gemm2_k128: bool, + do_finalize: bool = False, + routing_logits=None, +): + """The target's call, every argument explicit.""" + stack = _stack() + tokens = case["x"].shape[0] + if routing_logits is None: + # With the fused router only the shape is read; filled so that an + # autotuning run would profile on plausible routing. + routing_logits = ( + torch.randn(tokens, EXPERTS, device=DEV) if fused_router else case["logits"] + ) + return op( + routing_logits=routing_logits, + routing_bias=case["routing_bias"], + hidden_states=hidden_states, + hidden_states_scale=hidden_states_scale, + gemm1_weights=stack["gemm1_weights"], + gemm1_weights_scale=stack["gemm1_weights_scale"], + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, + gemm2_weights=stack["gemm2_weights"], + gemm2_weights_scale=stack["gemm2_weights_scale"], + gemm2_bias=None, + output1_scale_scalar=case["output1_scale_scalar"], + output1_scale_gate_scalar=case["output1_scale_gate_scalar"], + output2_scale_scalar=case["output2_scale_scalar"], + num_experts=EXPERTS, + top_k=TOP_K, + n_group=1, + topk_group=1, + intermediate_size=INTER, + local_expert_offset=0, + local_num_experts=EXPERTS, + routed_scaling_factor=ROUTED_SCALE, + routing_method_type=ROUTING_METHOD_DEEPSEEK_V3, + do_finalize=do_finalize, + act_type=ACT_TYPE_SWIGLU, + topk_weights=None, + topk_ids=None, + output=None, + tune_max_num_tokens=8192, + use_dp=False, + hidden_states_ready_event=None if event is None else event.cuda_event, + router_hidden_states=case["x"] if fused_router else None, + router_weight=case["router_weight"] if fused_router else None, + router_workspace=workspace if fused_router else None, + prefer_gemm2_k128=gemm2_k128, + ) + + +def _quantize_aside(case: dict, side: torch.cuda.Stream, delay: bool, out=None): + """The target's input path: quantize on `side`, forked after the current + stream's work, and return `(xq, xsf, event)` with `event` recorded there. + `out` holds buffers to write instead of fresh tensors; `delay` sleeps on + `side` first.""" + forked = torch.cuda.Event() + ready = torch.cuda.Event() + forked.record() + with torch.cuda.stream(side): + side.wait_event(forked) + if delay: + torch.cuda._sleep(SLEEP_CYCLES) + xq, xsf = _quantize(case["x"], case["g1"]) + if out is not None: + out[0].copy_(xq) + out[1].copy_(xsf) + xq, xsf = out + ready.record(side) + return xq, xsf, ready + + +def _valid_rows(outs): + """The unfinalized outputs as `(rows each slot reads, weights, map)`.""" + rows, weights, perm = outs + return rows[perm.long().reshape(-1)], weights, perm + + +def _finalize(outs) -> torch.Tensor: + """`sum_k weights[t, k] * rows[map[t, k]]` in fp32, rounded once: what the + post-MoE hop computes.""" + rows, weights, perm = outs + gathered = rows.float()[perm.long()] + return (gathered * weights.float().unsqueeze(-1)).sum(dim=1).to(torch.bfloat16) + + +def _assert_same(a, b, what: str) -> None: + for x, y, name in zip(_valid_rows(a), _valid_rows(b), ("rows", "weights", "map")): + assert torch.equal(x, y), f"{what}: {name} differ" + + +def _kernels(fn) -> list[str]: + """Names of the kernels `fn` launches, in the order they start.""" + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + fn() + torch.cuda.synchronize() + launched = [ + e + for e in prof.events() + if str(e.device_type) == "DeviceType.CUDA" and e.device_time_total > 0 + ] + return [e.name for e in sorted(launched, key=lambda e: e.time_range.start)] + + +def _k128_covers(tokens: int) -> bool: + """gemm2K128's gate at tile N 8: one CTA per SM, at most 15 units (token + tile x 128 output features) each, at most 64 token tiles.""" + tiles = min(tokens * TOP_K, EXPERTS) + sms = torch.cuda.get_device_properties(0).multi_processor_count + units = tiles * HIDDEN // 128 + return tiles <= 64 and -(-units // sms) <= 15 + + +# ── tests ───────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 + """Each cell as the target makes the call, against the reference. + + The late-join cells quantize on a side stream behind a sleep, so the + result also proves FC1 waited for the input. + """ + certifying(op) + spec = cell.spec + geometry = (spec["hidden"], spec["intermediate"], spec["num_experts"], spec["top_k"]) + assert geometry == (HIDDEN, INTER, EXPERTS, TOP_K), geometry + assert spec["routed_scaling_factor"] == ROUTED_SCALE and spec["local_expert_offset"] == 0 + case = _case(spec["tokens"], seed=zlib.crc32(cell.why.encode()) % 2**31, clear_margin=True) + workspace = make_fused_router_workspace(DEV) if spec["fused_router"] else None + if spec["late_join"]: + side = torch.cuda.Stream() + xq, xsf, ready = _quantize_aside(case, side, delay=True) + else: + xq, xsf, ready = case["xq"], case["xsf"], None + outs = _call( + case, + hidden_states=xq, + hidden_states_scale=xsf, + fused_router=spec["fused_router"], + workspace=workspace, + event=ready, + gemm2_k128=spec["gemm2_k128"], + ) + torch.cuda.synchronize() + + stack = _stack() + ref, _, ref_weights = op.reference( + routing_logits=None if spec["fused_router"] else case["logits"], + routing_bias=case["routing_bias"], + hidden_states=case["xq"], + hidden_states_scale=case["xsf"], + up=stack["up"], + gate=stack["gate"], + down=stack["down"], + output1_scale_scalar=case["output1_scale_scalar"], + output1_scale_gate_scalar=case["output1_scale_gate_scalar"], + output2_scale_scalar=case["output2_scale_scalar"], + top_k=TOP_K, + local_expert_offset=0, + routed_scaling_factor=ROUTED_SCALE, + router_hidden_states=case["x"] if spec["fused_router"] else None, + router_weight=case["router_weight"] if spec["fused_router"] else None, + ) + # Each token's weights, as a set: the kernel stores them in its top-k order. + torch.testing.assert_close( + outs[1].float().sort(dim=-1).values, + ref_weights.float().sort(dim=-1).values, + rtol=2.0**-7, + atol=0.0, + ) + _compare(_finalize(outs), ref) + + +@pytest.mark.parametrize("tokens", range(1, 9)) +def test_fused_router_matches_the_router_gemm(tokens) -> None: + """The fused router computes the logits of `dsv3_router_gemm_op` bit for + bit and routes them as the routing kernel does: every output is bitwise + the one the op gives when routed from those logits. Over every row count + the kernel has an instance for, on an off-grid router input where + summation order shows.""" + gen = torch.Generator(device=DEV).manual_seed(4000 + tokens) + case = _case(tokens, seed=4100 + tokens) + x = torch.randn(tokens, HIDDEN, device=DEV, generator=gen).to(torch.bfloat16) + weight = (torch.randn(EXPERTS, HIDDEN, device=DEV, generator=gen) * 0.02).to(torch.bfloat16) + case.update(x=x, router_weight=weight) + logits = torch.ops.trtllm.dsv3_router_gemm_op(x, weight.t(), None, torch.float32) + workspace = make_fused_router_workspace(DEV) + common = dict(hidden_states=case["xq"], hidden_states_scale=case["xsf"], gemm2_k128=False) + fused = _call(case, fused_router=True, workspace=workspace, **common) + routed = _call(case, fused_router=False, routing_logits=logits, **common) + _assert_same(fused, routed, f"{tokens} rows") + + +@pytest.mark.parametrize("tokens", (1, 8)) +def test_gemm2_k128_matches_the_cubin(tokens) -> None: + """FC2 at K = 128 writes the Gemm2 cubin's rows, bitwise, and runs where + its gate says it does.""" + case = _case(tokens, seed=4200 + tokens) + common = dict( + hidden_states=case["xq"], hidden_states_scale=case["xsf"], fused_router=False + ) + k128 = _call(case, gemm2_k128=True, **common) + cubin = _call(case, gemm2_k128=False, **common) + _assert_same(k128, cubin, f"{tokens} rows") + + names = _kernels(lambda: _call(case, gemm2_k128=True, **common)) + ran = any("gemm2K128Kernel" in n for n in names) + assert ran == _k128_covers(tokens), (names, _k128_covers(tokens)) + names = _kernels(lambda: _call(case, gemm2_k128=False, **common)) + assert not any("gemm2K128Kernel" in n for n in names), names + + +def test_c1_call_launches_three_kernels() -> None: + """The c1 call is the fused router, FC1 and FC2: no router GEMV, no + separate routing kernel.""" + case = _case(8, seed=4300) + workspace = make_fused_router_workspace(DEV) + shape_only = torch.randn(8, EXPERTS, device=DEV) + names = _kernels( + lambda: _call( + case, + hidden_states=case["xq"], + hidden_states_scale=case["xsf"], + fused_router=True, + workspace=workspace, + gemm2_k128=True, + routing_logits=shape_only, + ) + ) + assert len(names) == 3, names + assert "routingFusedRouterKernel" in names[0], names + assert ("gemm2K128Kernel" in names[2]) == _k128_covers(8), names + + +def test_late_join_waits_for_the_input() -> None: + """FC1 reads the input only after the event: the buffers handed to the op + hold a different activation's quantization until the side stream, held + back by a sleep, overwrites them. Routing launches without waiting.""" + case = _case(8, seed=4400) + stale = _case(8, seed=4401) + workspace = make_fused_router_workspace(DEV) + expected = _call( + case, + hidden_states=case["xq"], + hidden_states_scale=case["xsf"], + fused_router=True, + workspace=workspace, + gemm2_k128=True, + ) + buffers = (stale["xq"].clone(), stale["xsf"].clone()) + side = torch.cuda.Stream() + xq, xsf, ready = _quantize_aside(case, side, delay=True, out=buffers) + late = _call( + case, + hidden_states=xq, + hidden_states_scale=xsf, + fused_router=True, + workspace=workspace, + event=ready, + gemm2_k128=True, + ) + torch.cuda.synchronize() + _assert_same(late, expected, "late join") + + +def test_c1_call_replays_in_a_cuda_graph() -> None: + """The c1 step captured as the target runs it -- input quantized on a side + stream, joined inside the op -- replays the eager result on new inputs, and + every launch leaves the router workspace ready for the next one.""" + cases = [_case(8, seed=4500 + i) for i in range(3)] + static = {k: v.clone() for k, v in cases[0].items()} + workspace = make_fused_router_workspace(DEV) + side = torch.cuda.Stream() + + def step(): + xq, xsf, ready = _quantize_aside(static, side, delay=False) + return _call( + static, + hidden_states=xq, + hidden_states_scale=xsf, + fused_router=True, + workspace=workspace, + event=ready, + gemm2_k128=True, + routing_logits=static["logits"], + ) + + step() # load the kernels before capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = step() + + launches = 1 + for case in cases: + for key, value in case.items(): + static[key].copy_(value) + graph.replay() + launches += 1 + torch.cuda.synchronize() + eager = _call( + case, + hidden_states=case["xq"], + hidden_states_scale=case["xsf"], + fused_router=True, + workspace=workspace, + gemm2_k128=True, + ) + launches += 1 + torch.cuda.synchronize() + _assert_same(captured, eager, "replay") + words = workspace.view(torch.int32) + epoch = int(words[_EPOCH_WORD].item()) + assert epoch == launches, (epoch, launches) + current = words[(epoch & 1) * _BUFFER_WORDS : ((epoch & 1) + 1) * _BUFFER_WORDS] + assert bool((current == -(2**31)).all()), "the next launch's exchange buffer is not idle" + + +def test_fused_router_refuses_more_than_8_rows() -> None: + """Past 8 rows the fused router is a loud error; the caller routes from + the router GEMM's logits instead (the 32-row cell).""" + case = _case(9, seed=4600) + with pytest.raises(RuntimeError, match="fused router"): + _call( + case, + hidden_states=case["xq"], + hidden_states_scale=case["xsf"], + fused_router=True, + workspace=make_fused_router_workspace(DEV), + gemm2_k128=True, + ) + + +# ── guards ──────────────────────────────────────────────────────────────── + + +def test_guard_refuses_a_strided_router_input() -> None: + case = _case(2, seed=4700) + case["x"] = torch.cat([case["x"], case["x"]], dim=1)[:, ::2] + with pytest.raises(AssertionError, match="router_hidden_states must be contiguous"): + with validating(op): + _call( + case, + hidden_states=case["xq"], + hidden_states_scale=case["xsf"], + fused_router=True, + workspace=make_fused_router_workspace(DEV), + gemm2_k128=True, + ) + + +def test_guard_refuses_a_zero_filled_workspace() -> None: + """A zero-filled workspace reads as published zero sums: routing on zero + logits, with no error.""" + case = _case(2, seed=4701) + workspace = torch.zeros(FUSED_ROUTER_WORKSPACE_WORDS, dtype=torch.int32, device=DEV) + with pytest.raises(AssertionError, match="between-launch state"): + with validating(op): + _call( + case, + hidden_states=case["xq"], + hidden_states_scale=case["xsf"], + fused_router=True, + workspace=workspace, + gemm2_k128=True, + ) + + +def test_guard_refuses_an_unrecorded_event() -> None: + """An event never recorded has no cudaEvent_t yet: FC1 would not wait.""" + case = _case(2, seed=4702) + with pytest.raises(AssertionError, match="null cudaEvent_t"): + with validating(op): + _call( + case, + hidden_states=case["xq"], + hidden_states_scale=case["xsf"], + fused_router=True, + workspace=make_fused_router_workspace(DEV), + event=torch.cuda.Event(), + gemm2_k128=True, + ) From bf6f5597042d32f3c5cd3a1fac7033fdb4c4991e Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:12:48 -0700 Subject: [PATCH 18/53] Run the GLM-5.2 decode step's MoE blocks on the routed-experts and split-K ops Up to 8 rows, each MoE block runs its shared expert as two split-K GEMMs around SwiGLU on the aux stream and its routed experts as one trtllm-gen call on the main stream: the routing kernel computes the router GEMV, the NVFP4 input quantizes on its own stream and joins right before FC1, and FC2 runs on the K = 128 kernel. The decode target reads the model's dimensions off the pretrained config (the core has no config property). Signed-off-by: Dongfeng Yu --- .../glm_5_2_nvfp4__sm_107__tp16/modeling.py | 158 ++++++++++++++++-- 1 file changed, 147 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py index b926bb4f2939..8ad9e6836ad8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py @@ -19,6 +19,9 @@ import torch from transformers import PretrainedConfig +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata from tensorrt_llm._torch.attention.backends.sparse.dsa.module import ( forward_generation_sparse_attn, @@ -29,7 +32,14 @@ from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM from tensorrt_llm._torch.modules.gated_mlp import GatedMLP from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.modules.swiglu import swiglu +from tensorrt_llm._torch.moe.fused_moe.trtllm_gen.kernel_inputs import ( + extract_routing_params, + get_data_or_none, + to_trtllm_gen_act_type, +) from tensorrt_llm._torch.speculative import SpecMetadata +from tensorrt_llm._torch.utils import EventType from . import weights as _weights from .ops.embedding.fused_embedding import fused_embedding @@ -37,6 +47,12 @@ from .ops.gemm.fused_a_gemm import MAX_TOKENS as FUSED_A_MAX_TOKENS from .ops.gemm.fused_a_gemm import OUT_FEATURES as FUSED_A_OUT_FEATURES from .ops.gemm.fused_a_gemm import fused_a_gemm +from .ops.gemm.low_m_gemm import low_m_gemm +from .ops.moe.routed_experts import ( + FUSED_ROUTER_MAX_TOKENS, + make_fused_router_workspace, + routed_experts, +) # The deployment this target is written for; route() selects it on the same # facts, and the constructor states them again for a direct construction. @@ -91,6 +107,13 @@ class DecodeTarget: kv_a_proj_with_mqa weight (up to 8 rows), and the built-in modules compute the rest. + The MoE block over up to 8 rows runs its shared expert as two split-K + GEMMs around SwiGLU on the aux stream, and its routed experts in one + trtllm-gen call on the main stream: the routing kernel runs the router GEMV + itself, the NVFP4 input quantizes on its own stream and joins right before + FC1, and FC2 runs on the K = 128 kernel. A larger batch takes the built-in + MoE module. + The embedding table is whole on every rank, so the lookup is one local kernel. The dense MLP runs as two weight-streaming kernels on this rank's 768 intermediate columns for up to 8 rows (every c1 step: one request and @@ -99,13 +122,15 @@ class DecodeTarget: def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: self.core = core + config = core.model_config.pretrained_config + self.hidden_size = config.hidden_size for layer in core.layers[: core.num_hidden_layers]: mla = layer.self_attn assert layer.allreduce is not None and layer.next_layer_layernorm is not None assert not mla.apply_rotary_emb and not mla.mapping.has_cp_helix() assert tuple(mla.kv_a_proj_with_mqa.weight.shape) == ( FUSED_A_OUT_FEATURES, - core.config.hidden_size, + config.hidden_size, ) and mla.kv_a_proj_with_mqa.weight.dtype == torch.bfloat16, ( f"layer {layer.layer_idx}: kv_a_proj_with_mqa is not the BF16 [q_a; kv_a] weight" ) @@ -125,6 +150,41 @@ def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: and layer.fusion_config.POST_MOE_FUSION and layer.mapping.is_multi_node() ), f"layer {layer.layer_idx}: unexpected MoE-layer configuration" + shared = layer.mlp.shared_experts + assert ( + shared is not None + and layer.mlp.shared_output_scale is None + and shared._is_plain_swiglu() + and shared.swiglu_limit is None + and shared.gate_up_proj.weight.dtype == torch.bfloat16 + and shared.down_proj.weight.dtype == torch.bfloat16 + ), f"layer {layer.layer_idx}: the shared expert is not a plain BF16 SwiGLU MLP" + backend = getattr(layer.mlp.experts, "backend", layer.mlp.experts) + assert ( + backend.has_nvfp4 + and backend.scaling_vector_size == 16 + and getattr(backend, "fc31_act_scale", None) is None + and not backend.bias + and backend.w3_w1_weight.shape[-1] * 2 == config.hidden_size + ), f"layer {layer.layer_idx}: routed experts are not unpadded TRTLLM-Gen NVFP4" + moe_layers = [ + layer + for layer in core.layers[: core.num_hidden_layers] + if not isinstance(layer.mlp, GatedMLP) + ] + device = core.embed_tokens.weight.device + # One exchange workspace per router call site: its launches must run in stream order. + self._router_workspace = { + layer.layer_idx: make_fused_router_workspace(device) for layer in moe_layers + } + # The routing kernel runs the router GEMV, so the logits it is handed only give the + # shape; autotuner profiling routes from them, so they hold random values. + self._router_logits = torch.randn( + FUSED_ROUTER_MAX_TOKENS, config.n_routed_experts, device=device + ) + self._quant_stream = torch.cuda.Stream() + self._quant_fork = torch.cuda.Event() + self._quant_ready = torch.cuda.Event() def forward( self, @@ -142,7 +202,7 @@ def forward( hidden_states = core.layers[0].input_layernorm(inputs_embeds) rows = hidden_states.shape[0] small_m = torch.ops.trtllm.glm_dense_mlp_small_m_supported( - rows, core.config.hidden_size, INTER_PER_RANK + rows, self.hidden_size, INTER_PER_RANK ) fused_a = rows <= FUSED_A_MAX_TOKENS for layer in core.layers[: core.num_hidden_layers]: @@ -156,6 +216,80 @@ def forward( ) return hidden_states + def _experts(self, layer, hidden_states: torch.Tensor) -> torch.Tensor: + """The MoE block's output for up to 8 rows: shared expert plus routed experts.""" + mlp = layer.mlp + backend = getattr(mlp.experts, "backend", mlp.experts) + rows = hidden_states.shape[0] + + self._quant_fork.record() + with torch.cuda.stream(self._quant_stream): + self._quant_stream.wait_event(self._quant_fork) + x_fp4, x_sf = fp4_quantize( + hidden_states, backend.fc31_input_scale, backend.scaling_vector_size, False, False + ) + self._quant_ready.record() + routing = extract_routing_params(backend.routing_method) + + def routed() -> torch.Tensor: + return routed_experts( + self._router_logits[:rows], + routing.routing_bias, + x_fp4, + x_sf.flatten().view(torch.float8_e4m3fn), + backend.w3_w1_weight, + backend.w3_w1_weight_scale.view(torch.float8_e4m3fn), + None, + backend.act_alpha, + backend.act_beta, + backend.act_clamp, + backend.w2_weight, + backend.w2_weight_scale.view(torch.float8_e4m3fn), + None, + get_data_or_none(backend, "fc31_scale_c"), + get_data_or_none(backend, "fc31_alpha"), + get_data_or_none(backend, "fc2_alpha"), + backend.num_slots, + routing.top_k, + routing.n_group, + routing.topk_group, + backend.w3_w1_weight.shape[-2] // 2, + backend.slot_start, + backend.expert_size_per_partition, + routing.routed_scaling_factor, + backend.routing_method.routing_method_type, + True, + to_trtllm_gen_act_type(backend.activation_type), + None, + None, + None, + backend.max_num_tokens, + False, + self._quant_ready.cuda_event, + hidden_states, + mlp.gate.weight, + self._router_workspace[layer.layer_idx], + True, + )[0] + + def shared() -> torch.Tensor: + expert = mlp.shared_experts + gate_up = expert.gate_up_proj.weight + h = low_m_gemm(hidden_states, gate_up, hidden_states.new_empty(rows, gate_up.shape[0])) + h = swiglu(h) + return low_m_gemm( + h, expert.down_proj.weight, hidden_states.new_empty(rows, hidden_states.shape[1]) + ) + + routed_output, shared_output = maybe_execute_in_parallel( + routed, + shared, + mlp.event_dict[EventType.Main], + mlp.event_dict[EventType.MoeShared], + mlp.aux_stream, + ) + return shared_output.add_(routed_output) + @staticmethod def _attention(layer, position_ids, hidden_states, attn_metadata, all_reduce_params, fused_a): """MLA with the DSA indexer over generation rows, through o_proj.""" @@ -233,8 +367,7 @@ def _dense_layer( ), ) - @staticmethod - def _moe_layer(layer, position_ids, hidden_states, residual, attn_metadata, fused_a): + def _moe_layer(self, layer, position_ids, hidden_states, residual, attn_metadata, fused_a): hidden_states = DecodeTarget._attention( layer, position_ids, @@ -253,13 +386,16 @@ def _moe_layer(layer, position_ids, hidden_states, residual, attn_metadata, fuse trigger_completion_at_end=False, ), ) - hidden_states = layer.mlp( - hidden_states, - None, - all_rank_num_tokens=attn_metadata.all_rank_num_tokens, - final_all_reduce_params=AllReduceParams(enable_allreduce=False), - do_finalize=True, - ) + if hidden_states.shape[0] <= FUSED_ROUTER_MAX_TOKENS: + hidden_states = self._experts(layer, hidden_states) + else: + hidden_states = layer.mlp( + hidden_states, + None, + all_rank_num_tokens=attn_metadata.all_rank_num_tokens, + final_all_reduce_params=AllReduceParams(enable_allreduce=False), + do_finalize=True, + ) return layer.allreduce( hidden_states, all_reduce_params=AllReduceParams( From 14a26accb3428e3ec63400932916c9300e0874c7 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:08:00 -0700 Subject: [PATCH 19/53] Add the GLM-5.2 tp16 target's fused MLA KV-half append op glm_mla_kv_half_append (kernels/glmMoeDsa/glmMlaKvHalfAppend.cu, thop/glmMlaKvHalfAppendOp.cpp) runs the KV half of an MLA generation step in one kernel: RMSNorm of the 512 latent columns of the raw kv_a_proj output, GPT-J RoPE of the 64 K-RoPE columns, E4M3 quantization at unit scale straight from fp32, and the append to the paged FP8 cache through the block offsets, pool pointers and pool mapping that mla_rope_generation takes. Block 0 also writes the generation attention's scales. With trigger_at_entry its PDL dependents may launch at its entry; they read its outputs after their own grid-dependency wait. The modeling_v2 entry ops/attention/mla_kv_half_append.py states the op, its reference and its cells: the verify and MTP draft step 0 rows of both fused-A widths, a draft-step row and two sequences. Its sm_107 test reads the appended rows back where convert_req_index_to_global places their positions on a DSACacheManagerV2 pool. Signed-off-by: Dongfeng Yu --- .../kernels/glmMoeDsa/glmMlaKvHalfAppend.cu | 231 ++++++++++++++ .../kernels/glmMoeDsa/glmMlaKvHalfAppend.h | 56 ++++ cpp/tensorrt_llm/thop/CMakeLists.txt | 3 +- .../thop/glmMlaKvHalfAppendOp.cpp | 148 +++++++++ .../ops/attention/mla_kv_half_append.py | 294 ++++++++++++++++++ .../test_modeling_v2_mla_kv_half_append.py | 290 +++++++++++++++++ 6 files changed, 1021 insertions(+), 1 deletion(-) create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.cu create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.h create mode 100644 cpp/tensorrt_llm/thop/glmMlaKvHalfAppendOp.cpp create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/mla_kv_half_append.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_mla_kv_half_append.py diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.cu new file mode 100644 index 000000000000..43513de09ad2 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.cu @@ -0,0 +1,231 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ +namespace +{ + +constexpr int kNopeDim = kGlmKvHalfLatentDim; +constexpr int kRopeDim = kGlmKvHalfRopeDim; +// BF16 values per 16-byte vector. +constexpr int kEltsPerVec = 8; +constexpr int kNopeVecs = kNopeDim / kEltsPerVec; +constexpr int kRopeVecs = kRopeDim / kEltsPerVec; +// The cache row in 8-element units. +constexpr int kRowVecs = kNopeVecs + kRopeVecs; +constexpr int kWarpSize = 32; +constexpr int kNopeVecsPerLane = kNopeVecs / kWarpSize; +// Four token rows per block, one warp per row. +constexpr int kBlockSize = 128; +static_assert(kNopeVecs % kWarpSize == 0 && kRopeVecs <= kWarpSize, "the row splits over one warp"); + +inline __device__ float2 bf16x2ToFloat2(uint32_t word) +{ + return make_float2(__uint_as_float(word << 16), __uint_as_float(word & 0xFFFF0000u)); +} + +// Eight FP32 values to saturated E4M3. +inline __device__ uint2 toFp8x8(float const (&v)[kEltsPerVec]) +{ + uint2 out; + auto* pairs = reinterpret_cast<__nv_fp8x2_storage_t*>(&out); +#pragma unroll + for (int j = 0; j < kEltsPerVec / 2; ++j) + { + pairs[j] = __nv_cvt_float2_to_fp8x2(make_float2(v[2 * j], v[2 * j + 1]), __NV_SATFINITE, __NV_E4M3); + } + return out; +} + +// One warp per token row. Lane l holds latent vectors l and l + 32 and, for l < 8, K-RoPE vector l. +// kTriggerAtEntry: the PDL dependents may launch at kernel entry instead of after the stores; they read this kernel's +// outputs only after their own grid-dependency wait. +template +__global__ void __launch_bounds__(kBlockThreads) glmMlaKvHalfAppendKernel(__nv_bfloat16 const* __restrict__ latent, + int latentRowStride, KVBlockArray kvCache, float2 const* __restrict__ cosSinCache, + __nv_bfloat16 const* __restrict__ kvNormWeight, float kvNormEps, int numTokens, int seqLenQ, + int32_t const* __restrict__ kvLens, float* __restrict__ bmm1Scale, float* __restrict__ bmm2Scale, + float softmaxScale) +{ + constexpr int kRowsPerBlock = kBlockThreads / kWarpSize; + int const lane = threadIdx.x % kWarpSize; + int const warpId = threadIdx.x / kWarpSize; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if constexpr (kTriggerAtEntry) + { + cudaTriggerProgrammaticLaunchCompletion(); + } + // The latent row is the previous kernel's output: wait before any global read. + cudaGridDependencySynchronize(); +#endif + + if (blockIdx.x == 0 && threadIdx.x == 0) + { + // The scales of the generation attention over this cache (unit KV scales). + constexpr float kLog2e = 1.4426950408889634074f; + bmm1Scale[0] = softmaxScale; + bmm1Scale[1] = softmaxScale * kLog2e; + bmm2Scale[0] = 1.f; + } + + // The norm weight is row-invariant: two 8-element slices per lane, kept in FP32. + float weight[kNopeVecsPerLane][kEltsPerVec]; +#pragma unroll + for (int i = 0; i < kNopeVecsPerLane; ++i) + { + uint4 const w = *reinterpret_cast(kvNormWeight + (lane + i * kWarpSize) * kEltsPerVec); + uint32_t const words[4] = {w.x, w.y, w.z, w.w}; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float2 const f = bf16x2ToFloat2(words[j]); + weight[i][2 * j] = f.x; + weight[i][2 * j + 1] = f.y; + } + } + + for (int token = warpId + blockIdx.x * kRowsPerBlock; token < numTokens; token += kRowsPerBlock * gridDim.x) + { + int const seq = token / seqLenQ; + int const tokenInSeq = token - seq * seqLenQ; + // Cache slot and RoPE position, as the generation RoPE kernel derives them. + int const tokenKvIdx = kvLens[seq] - seqLenQ + tokenInSeq; + if (tokenKvIdx < 0) + { + continue; + } + __nv_bfloat16 const* row = latent + static_cast(token) * static_cast(latentRowStride); + // Page-table read first: it does not depend on the reduction. + auto* kDst = reinterpret_cast(kvCache.getKBlockPtr(seq, tokenKvIdx)); + + uint4 nope[kNopeVecsPerLane]; +#pragma unroll + for (int i = 0; i < kNopeVecsPerLane; ++i) + { + nope[i] = *reinterpret_cast(row + (lane + i * kWarpSize) * kEltsPerVec); + } + uint4 rope = make_uint4(0, 0, 0, 0); + if (lane < kRopeVecs) + { + rope = *reinterpret_cast(row + kNopeDim + lane * kEltsPerVec); + } + + float x[kNopeVecsPerLane][kEltsPerVec]; + float sumSquares = 0.f; +#pragma unroll + for (int i = 0; i < kNopeVecsPerLane; ++i) + { + uint32_t const words[4] = {nope[i].x, nope[i].y, nope[i].z, nope[i].w}; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float2 const f = bf16x2ToFloat2(words[j]); + x[i][2 * j] = f.x; + x[i][2 * j + 1] = f.y; + sumSquares += f.x * f.x + f.y * f.y; + } + } +#pragma unroll + for (int mask = kWarpSize / 2; mask > 0; mask >>= 1) + { + sumSquares += __shfl_xor_sync(0xFFFFFFFFu, sumSquares, mask); + } + float const rstd = rsqrtf(sumSquares / static_cast(kNopeDim) + kvNormEps); + +#pragma unroll + for (int i = 0; i < kNopeVecsPerLane; ++i) + { + float scaled[kEltsPerVec]; +#pragma unroll + for (int j = 0; j < kEltsPerVec; ++j) + { + scaled[j] = x[i][j] * (rstd * weight[i][j]); + } + int const vec = lane + i * kWarpSize; + *reinterpret_cast(kDst + kvCache.getKVLocalIdx(tokenKvIdx, 0, kRowVecs, vec) * kEltsPerVec) + = toFp8x8(scaled); + } + + if (lane < kRopeVecs) + { + // Pair p of the RoPE slice uses entry kRopeDim * position + p of the table. + float2 const* coef = cosSinCache + static_cast(kRopeDim) * tokenKvIdx + lane * (kEltsPerVec / 2); + uint32_t const words[4] = {rope.x, rope.y, rope.z, rope.w}; + float rotated[kEltsPerVec]; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float2 const v = bf16x2ToFloat2(words[j]); + float2 const cs = coef[j]; + rotated[2 * j] = cs.x * v.x - cs.y * v.y; + rotated[2 * j + 1] = cs.x * v.y + cs.y * v.x; + } + *reinterpret_cast( + kDst + kvCache.getKVLocalIdx(tokenKvIdx, 0, kRowVecs, kNopeVecs + lane) * kEltsPerVec) + = toFp8x8(rotated); + } + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if constexpr (!kTriggerAtEntry) + { + cudaTriggerProgrammaticLaunchCompletion(); + } +#endif +} + +} // namespace + +void invokeGlmMlaKvHalfAppend(__nv_bfloat16 const* latent, int32_t latentRowStride, KVBlockArray const& kvCache, + float2 const* cosSinCache, __nv_bfloat16 const* kvNormWeight, float kvNormEps, int32_t numTokens, int32_t seqLenQ, + int32_t const* kvLens, float* bmm1Scale, float* bmm2Scale, float softmaxScale, bool triggerAtEntry, + cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(latent != nullptr && cosSinCache != nullptr && kvNormWeight != nullptr && kvLens != nullptr + && bmm1Scale != nullptr && bmm2Scale != nullptr, + "The GLM MLA KV append needs the latent rows, the rope table, the norm weight, the KV lengths and the scales"); + TLLM_CHECK_WITH_INFO(latentRowStride >= kNopeDim + kRopeDim && latentRowStride % kEltsPerVec == 0, + "latent row stride %d must cover a %d-column row and be a multiple of %d", latentRowStride, kNopeDim + kRopeDim, + kEltsPerVec); + TLLM_CHECK_WITH_INFO(numTokens >= 0 && seqLenQ > 0 && numTokens % seqLenQ == 0, + "MLA can only support input sequences with the same sequence length."); + + constexpr int kRowsPerBlock = kBlockSize / kWarpSize; + dim3 const grid(static_cast(tensorrt_llm::common::divUp(std::max(numTokens, 1), kRowsPerBlock))); + auto const kernel = triggerAtEntry ? glmMlaKvHalfAppendKernel + : glmMlaKvHalfAppendKernel; + tensorrt_llm::common::launchWithPdlWhenEnabled("glmMlaKvHalfAppend", kernel, grid, kBlockSize, 0, stream, latent, + static_cast(latentRowStride), kvCache, cosSinCache, kvNormWeight, kvNormEps, static_cast(numTokens), + static_cast(seqLenQ), kvLens, bmm1Scale, bmm2Scale, softmaxScale); +} + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.h new file mode 100644 index 000000000000..219dd77ad8f0 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.h @@ -0,0 +1,56 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/kernels/kvCacheUtils.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +// The MLA latent row of GLM-5.2: kv_lora_rank 512 latent columns followed by qk_rope_head_dim 64 K-RoPE columns. +constexpr int32_t kGlmKvHalfLatentDim = 512; +constexpr int32_t kGlmKvHalfRopeDim = 64; + +// The KV half of MLA generation for the latent layout above, in one kernel: RMSNorm of the 512 latent columns of the +// raw kv_a_proj output, GPT-J RoPE of the 64 K-RoPE columns, E4M3 quantization at unit scale straight from FP32, and +// the append to the paged FP8 cache through kvCache, at the slot and RoPE position of the generation RoPE kernel +// (position kvLens[s] - seqLenQ + t for token t of sequence s; a negative position is skipped). Block 0 also writes +// the generation attention's scales: bmm1Scale = {softmaxScale, softmaxScale * log2(e)}, bmm2Scale = {1}. +// - latent: [numTokens] rows of at least 576 BF16 values, latentRowStride elements apart (a multiple of 8, rows +// 16-byte aligned); +// - cosSinCache: per position, 64 (cos, sin) float pairs (the duplicated-data table of the MLA rope), of which pair +// p < 32 rotates K-RoPE columns 2p and 2p + 1; +// - kvNormWeight: the 512 kv_a_layernorm weights. +// PDL: the kernel waits on its grid dependency before any global read. With triggerAtEntry its dependents may launch +// at its entry, otherwise after its stores; either way they read its outputs only after their own grid-dependency +// wait. +void invokeGlmMlaKvHalfAppend(__nv_bfloat16 const* latent, int32_t latentRowStride, KVBlockArray const& kvCache, + float2 const* cosSinCache, __nv_bfloat16 const* kvNormWeight, float kvNormEps, int32_t numTokens, int32_t seqLenQ, + int32_t const* kvLens, float* bmm1Scale, float* bmm2Scale, float softmaxScale, bool triggerAtEntry, + cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 4ebb89865614..4f925e421592 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -74,7 +74,8 @@ add_th_common_family( dsv3RopeOp.cpp convertReqIndexToGlobalOp.cpp nvfp4MlaKvCacheGatherOp.cpp - trtllmGenQKVProcessOp.cpp) + trtllmGenQKVProcessOp.cpp + glmMlaKvHalfAppendOp.cpp) add_th_common_family( th_common_gemm_quant diff --git a/cpp/tensorrt_llm/thop/glmMlaKvHalfAppendOp.cpp b/cpp/tensorrt_llm/thop/glmMlaKvHalfAppendOp.cpp new file mode 100644 index 000000000000..3726c3f61459 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmMlaKvHalfAppendOp.cpp @@ -0,0 +1,148 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/quantization.h" +#include "tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.h" +#include "tensorrt_llm/runtime/torchUtils.h" +#include "tensorrt_llm/thop/attentionOp.h" + +#include +#include + +#include +#include +#include + +namespace th = torch; +namespace tc = tensorrt_llm::common; +namespace tk = tensorrt_llm::kernels::glm_moe_dsa; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kInt32Max = std::numeric_limits::max(); +constexpr int64_t kLatentRowDim = tk::kGlmKvHalfLatentDim + tk::kGlmKvHalfRopeDim; +// The kernel walks the latent rows with 16-byte vector loads. +constexpr int64_t kVectorBytes = 16; +// Per position, the rope table holds 64 (cos, sin) float pairs. +constexpr int64_t kRopeTableFloatsPerPosition = 2 * tk::kGlmKvHalfRopeDim; + +} // namespace + +// The KV half of GLM-5.2's MLA generation step: kv_a_layernorm of the 512 latent columns of the raw kv_a_proj output, +// RoPE of its 64 K columns, E4M3 at unit scale, appended to layer layer_idx of the paged FP8 cache at the generation +// positions sequence_length[num_contexts + s] - seq_len_q + t, plus the generation attention's scales in +// mla_bmm1_scale / mla_bmm2_scale. The paged-cache arguments are those of mla_rope_generation. +void glmMlaKvHalfAppend(th::Tensor const& latent, th::Tensor const& kvNormWeight, double kvNormEps, + th::Tensor const& rotaryCosSin, th::Tensor const& sequenceLength, int64_t numContexts, int64_t seqLenQ, + th::Tensor const& kvCacheBlockOffsets, th::Tensor const& hostKvCachePoolPointers, + th::Tensor const& hostKvCachePoolMapping, int64_t layerIdx, int64_t tokensPerBlock, int64_t attentionWindowSize, + th::Tensor& mlaBmm1Scale, th::Tensor& mlaBmm2Scale, double softmaxScale, bool triggerAtEntry) +{ + TORCH_CHECK(latent.is_cuda() && kvNormWeight.is_cuda() && rotaryCosSin.is_cuda() && sequenceLength.is_cuda() + && kvCacheBlockOffsets.is_cuda() && mlaBmm1Scale.is_cuda() && mlaBmm2Scale.is_cuda(), + "glm_mla_kv_half_append takes CUDA tensors besides the host pool pointers and mapping"); + auto const device = latent.device(); + TORCH_CHECK(kvNormWeight.device() == device && rotaryCosSin.device() == device && sequenceLength.device() == device + && kvCacheBlockOffsets.device() == device && mlaBmm1Scale.device() == device + && mlaBmm2Scale.device() == device, + "glm_mla_kv_half_append takes its CUDA tensors on one device"); + + TORCH_CHECK(latent.scalar_type() == th::kBFloat16 && latent.dim() == 2 && latent.size(1) >= kLatentRowDim + && latent.stride(1) == 1, + "latent must be a BF16 [tokens, >= ", kLatentRowDim, "] view with unit column stride"); + TORCH_CHECK(latent.stride(0) % (kVectorBytes / 2) == 0 && latent.stride(0) <= kInt32Max + && reinterpret_cast(latent.data_ptr()) % kVectorBytes == 0, + "latent rows must be 16-byte aligned (row stride a multiple of 8 elements), got row stride ", latent.stride(0)); + int64_t const numTokens = latent.size(0); + TORCH_CHECK(numTokens <= kInt32Max, "too many latent rows"); + TORCH_CHECK(seqLenQ > 0 && numTokens % seqLenQ == 0, + "seq_len_q must be positive and divide the latent rows; got ", seqLenQ, " for ", numTokens, " rows"); + int64_t const numGenerations = numTokens / seqLenQ; + + TORCH_CHECK(kvNormWeight.scalar_type() == th::kBFloat16 && kvNormWeight.is_contiguous() + && kvNormWeight.numel() == tk::kGlmKvHalfLatentDim + && reinterpret_cast(kvNormWeight.data_ptr()) % kVectorBytes == 0, + "kv_norm_weight must be a contiguous, 16-byte aligned BF16 tensor of ", tk::kGlmKvHalfLatentDim, " elements"); + TORCH_CHECK(rotaryCosSin.scalar_type() == th::kFloat32 && rotaryCosSin.is_contiguous() + && rotaryCosSin.numel() % kRopeTableFloatsPerPosition == 0, + "rotary_cos_sin must be a contiguous float32 table of ", kRopeTableFloatsPerPosition, " values per position"); + TORCH_CHECK(sequenceLength.scalar_type() == th::kInt32 && sequenceLength.dim() == 1 + && sequenceLength.is_contiguous() && numContexts >= 0 + && sequenceLength.size(0) >= numContexts + numGenerations, + "sequence_length must be a contiguous int32 tensor covering the context and generation sequences"); + // KVBlockArray addresses sequence s of the pool's rows at s * 2 * max_blocks. + TORCH_CHECK(kvCacheBlockOffsets.scalar_type() == th::kInt32 && kvCacheBlockOffsets.dim() == 4 + && kvCacheBlockOffsets.size(2) == 2 && kvCacheBlockOffsets.stride(3) == 1 + && kvCacheBlockOffsets.stride(2) == kvCacheBlockOffsets.size(3) + && kvCacheBlockOffsets.stride(1) == 2 * kvCacheBlockOffsets.size(3) + && kvCacheBlockOffsets.size(1) >= numContexts + numGenerations, + "kv_cache_block_offsets must be an int32 [pools, sequences, 2, max_blocks] tensor, dense in its last three " + "dimensions"); + TORCH_CHECK(hostKvCachePoolPointers.device().is_cpu() && hostKvCachePoolMapping.device().is_cpu(), + "host_kv_cache_pool_pointers and host_kv_cache_pool_mapping must be CPU tensors"); + TORCH_CHECK(tokensPerBlock > 0 && (tokensPerBlock & (tokensPerBlock - 1)) == 0, + "tokens_per_block must be a power of two, got ", tokensPerBlock); + TORCH_CHECK(attentionWindowSize > 0 && attentionWindowSize <= kInt32Max, "invalid attention_window_size ", + attentionWindowSize); + TORCH_CHECK(mlaBmm1Scale.scalar_type() == th::kFloat32 && mlaBmm1Scale.is_contiguous() && mlaBmm1Scale.numel() >= 2, + "mla_bmm1_scale must hold two contiguous float32 values"); + TORCH_CHECK(mlaBmm2Scale.scalar_type() == th::kFloat32 && mlaBmm2Scale.numel() >= 1, + "mla_bmm2_scale must hold a float32 value"); + + // The latent pool of the FP8 MLA cache: one 576-element head per token, one byte per element. + auto const kvCacheBuffers = buildPagedKvCacheBuffers(std::optional(kvCacheBlockOffsets), + std::optional(hostKvCachePoolPointers), std::optional(hostKvCachePoolMapping), + tc::QuantMode::fp8KvCache(), layerIdx, numGenerations, tokensPerBlock, /*kv_head_num=*/1, kLatentRowDim, + attentionWindowSize, attentionWindowSize, /*beam_width=*/1, numContexts, /*is_mla_enable=*/true, + latent.element_size()); + TORCH_CHECK(kvCacheBuffers.kvCacheBuffer.data != nullptr, "glm_mla_kv_half_append found no paged KV cache"); + + auto const stream = at::cuda::getCurrentCUDAStream(latent.get_device()); + tk::invokeGlmMlaKvHalfAppend(static_cast<__nv_bfloat16 const*>(latent.data_ptr()), + static_cast(latent.stride(0)), kvCacheBuffers.kvCacheBuffer, + reinterpret_cast(rotaryCosSin.data_ptr()), + static_cast<__nv_bfloat16 const*>(kvNormWeight.data_ptr()), static_cast(kvNormEps), + static_cast(numTokens), static_cast(seqLenQ), + sequenceLength.data_ptr() + numContexts, mlaBmm1Scale.data_ptr(), + mlaBmm2Scale.data_ptr(), static_cast(softmaxScale), triggerAtEntry, stream); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_mla_kv_half_append(Tensor latent, Tensor kv_norm_weight, float kv_norm_eps, Tensor rotary_cos_sin, " + "Tensor sequence_length, int num_contexts, int seq_len_q, Tensor kv_cache_block_offsets, " + "Tensor host_kv_cache_pool_pointers, Tensor host_kv_cache_pool_mapping, int layer_idx, int tokens_per_block, " + "int attention_window_size, Tensor(a!) mla_bmm1_scale, Tensor(b!) mla_bmm2_scale, float softmax_scale, " + "bool trigger_at_entry=False) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_mla_kv_half_append", &tensorrt_llm::torch_ext::glmMlaKvHalfAppend); +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/mla_kv_half_append.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/mla_kv_half_append.py new file mode 100644 index 000000000000..d777072dfc1b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/mla_kv_half_append.py @@ -0,0 +1,294 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The KV half of an MLA generation step: kv_a RMSNorm, K RoPE and the FP8 paged-cache append.""" + +from __future__ import annotations + +import math + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The latent row: kv_lora_rank 512 normalized columns followed by qk_rope_head_dim 64 K-RoPE +#: columns. +KV_LORA_RANK = 512 +ROPE_DIM = 64 +#: Where the latent starts in fused-A's output row (after q_lora_rank 2048 q_a columns). +LATENT_OFFSET = 2048 +#: The E4M3 range the kernel saturates to. +_E4M3_MAX = 448.0 +#: log2(e) as the kernel's fp32 constant. +_LOG2E = 1.4426950408889634 + +#: GLM-5.2's pool page, kv_a_layernorm epsilon, 1 / sqrt(qk_nope_head_dim + qk_rope_head_dim) +#: and rope_theta. +_GLM = dict( + tokens_per_block=64, + kv_norm_eps=1e-5, + softmax_scale=1.0 / math.sqrt(256.0), + rope_theta=8_000_000.0, +) + + +class GlmMlaKvHalfAppend(OpWrapper): + """One generation step's MLA KV half over a paged FP8 latent cache. + + Wraps ``torch.ops.trtllm.glm_mla_kv_half_append``. Row ``t`` of ``latent`` is token + ``t % seq_len_q`` of generation sequence ``s = t // seq_len_q`` and sits at position + ``p = sequence_length[num_contexts + s] - seq_len_q + t % seq_len_q``. Its cache row is:: + + e4m3(rmsnorm(latent[t, :512], kv_norm_weight, kv_norm_eps)) # x * (rstd * w), fp32 + e4m3(gptj_rope(latent[t, 512:576], position p)) # pairs (2i, 2i + 1), fp32 + + saturated to +-448 at unit scale, written to layer ``layer_idx`` of the paged pool at + position ``p`` through the same block offsets, pool pointers and pool mapping as + ``mla_rope_generation``. A negative position is skipped. The call also fills + ``mla_bmm1_scale = [softmax_scale, softmax_scale * log2(e)]`` and ``mla_bmm2_scale = [1]``, + the scales the generation attention reads. It returns nothing. + + Values reach E4M3 straight from fp32: unlike ``mla_rope_generation``'s KV half, neither + the normalized latent nor the rotated K is rounded to BF16 first. + + Fusion boundary: the RMSNorm of the latent, the RoPE of K, the quantization, the append + and the scales are inside. kv_a_proj, the query and the attention are not. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="verify and MTP step 0: 8 rows of fused-A's 2624-column output across a page edge", + spec=dict( + seq_len_q=8, + cached=(4_090,), + latent_width=2624, + num_layers=4, + layer_idx=1, + trigger_at_entry=True, + **_GLM, + ), + ), + Cell( + why="full-indexer layers: the latent inside fused-A's 2784-column output (wk|wp rows)", + spec=dict( + seq_len_q=8, + cached=(4_090,), + latent_width=2784, + num_layers=4, + layer_idx=0, + trigger_at_entry=True, + **_GLM, + ), + ), + Cell( + why="MTP draft steps 1-6: one row at the MTP layer, the last layer of the pool", + spec=dict( + seq_len_q=1, + cached=(4_100,), + latent_width=2624, + num_layers=4, + layer_idx=3, + trigger_at_entry=True, + **_GLM, + ), + ), + Cell( + why="two sequences (max_batch_size 2): per-sequence offsets and positions, end trigger", + spec=dict( + seq_len_q=8, + cached=(700, 4_090), + latent_width=2624, + num_layers=4, + layer_idx=2, + trigger_at_entry=False, + **_GLM, + ), + ), + ) + + note = """ + Metadata consumed: none -- every input is an argument. What the target passes: + - `latent`: fused-A's output columns [2048, 2624) -- the raw kv_a_proj row, not + normalized; `kv_norm_weight` / `kv_norm_eps` the kv_a_layernorm weight and epsilon. + - `rotary_cos_sin`: the MLA rope table (`RopeParams(..., duplicate_data=True)`, GPT-J + pairs), covering every position the cache can hold (`max_seq_len`). + - `sequence_length`, `num_contexts`, `kv_cache_block_offsets`: the attention metadata's + `kv_lens_cuda_runtime`, `num_contexts` and `kv_cache_block_offsets`; + `host_kv_cache_pool_pointers` / `host_kv_cache_pool_mapping` the cache manager's + `kv_cache_pool_pointers` / `kv_cache_pool_mapping`; `layer_idx` the local layer; + `attention_window_size` the metadata's `max_seq_len`. + - `mla_bmm1_scale` ([2] fp32) and `mla_bmm2_scale` ([1] fp32): buffers allocated once, + read by `glm_sparse_mla_flash_decode`. + + The cache carries no scaling factor: GLM-5.2's FP8 cache is written and read at 1.0, as + the PyTorch MLA path writes it on the context server. The pool must hold plain 576-byte + E4M3 rows (`kv_cache_dtype: fp8`); the inline-scale `fp8_ds_mla` layout has 656-byte rows, + which the paged-cache arguments do not reveal, so the target asserts the KV dtype at + construction. + + PDL: the kernel waits on its grid dependency before any global read. With + `trigger_at_entry` its dependents may launch at its entry (the flash decode that gathers + earlier steps' rows before its wait), otherwise after its stores; every dependent must + wait on its own grid dependency before reading the appended rows or the scales. + + Not certified: a mixed batch (`num_contexts` > 0), more than two sequences, a ragged batch, + pages other than 64 tokens. Rewriting a prepared step rewrites the same slots. + """ + + def __call__( + self, + latent: torch.Tensor, + kv_norm_weight: torch.Tensor, + kv_norm_eps: float, + rotary_cos_sin: torch.Tensor, + sequence_length: torch.Tensor, + num_contexts: int, + seq_len_q: int, + kv_cache_block_offsets: torch.Tensor, + host_kv_cache_pool_pointers: torch.Tensor, + host_kv_cache_pool_mapping: torch.Tensor, + layer_idx: int, + tokens_per_block: int, + attention_window_size: int, + mla_bmm1_scale: torch.Tensor, + mla_bmm2_scale: torch.Tensor, + softmax_scale: float, + trigger_at_entry: bool = False, + ) -> None: + torch.ops.trtllm.glm_mla_kv_half_append( + latent, + kv_norm_weight, + kv_norm_eps, + rotary_cos_sin, + sequence_length, + num_contexts, + seq_len_q, + kv_cache_block_offsets, + host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping, + layer_idx, + tokens_per_block, + attention_window_size, + mla_bmm1_scale, + mla_bmm2_scale, + softmax_scale, + trigger_at_entry, + ) + + def reference( + self, + latent: torch.Tensor, + kv_norm_weight: torch.Tensor, + kv_norm_eps: float, + rotary_cos_sin: torch.Tensor, + sequence_length: torch.Tensor, + num_contexts: int, + seq_len_q: int, + kv_cache_block_offsets: torch.Tensor, + host_kv_cache_pool_pointers: torch.Tensor, + host_kv_cache_pool_mapping: torch.Tensor, + layer_idx: int, + tokens_per_block: int, + attention_window_size: int, + mla_bmm1_scale: torch.Tensor, + mla_bmm2_scale: torch.Tensor, + softmax_scale: float, + trigger_at_entry: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """``(rows, positions, bmm1, bmm2)``: what the op is supposed to write, and where. + + ``rows`` [T, 576] E4M3 is the cache row of each latent row, ``positions`` [T] int64 its + position in its sequence. The paged-cache arguments (block offsets, pool pointers and + mapping, layer, page size, window) only say where a row lands, and the test reads it + back there; ``mla_bmm1_scale`` / ``mla_bmm2_scale`` are where the scales land and + ``trigger_at_entry`` when the dependents launch, so all of them are ignored here. + + fp32 throughout, as the kernel: the RMSNorm multiplies ``x * (rstd * w)``, the RoPE + pairs consecutive columns (GPT-J), and each value is saturated and rounded to E4M3 + once. + """ + num_tokens = latent.shape[0] + device = latent.device + token = torch.arange(num_tokens, device=device) + seq = token // seq_len_q + positions = sequence_length[num_contexts + seq].long() - seq_len_q + token % seq_len_q + + x = latent[:, :KV_LORA_RANK].float() + rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + kv_norm_eps) + nope = x * (rstd * kv_norm_weight.float()) + + pairs = latent[:, KV_LORA_RANK : KV_LORA_RANK + ROPE_DIM].float().reshape( + num_tokens, ROPE_DIM // 2, 2 + ) + table = rotary_cos_sin.reshape(-1, ROPE_DIM, 2)[positions.clamp(min=0), : ROPE_DIM // 2] + cos, sin = table[..., 0], table[..., 1] + roped = torch.stack( + (cos * pairs[..., 0] - sin * pairs[..., 1], cos * pairs[..., 1] + sin * pairs[..., 0]), + dim=-1, + ).reshape(num_tokens, ROPE_DIM) + + rows = ( + torch.cat((nope, roped), dim=-1) + .clamp(-_E4M3_MAX, _E4M3_MAX) + .to(torch.float8_e4m3fn) + ) + scale = torch.tensor(softmax_scale, dtype=torch.float32, device=device) + log2e = torch.tensor(_LOG2E, dtype=torch.float32, device=device) + bmm1 = torch.stack((scale, scale * log2e)) + bmm2 = torch.ones(1, dtype=torch.float32, device=device) + return rows, positions, bmm1, bmm2 + + def is_valid( + self, + latent: torch.Tensor, + kv_norm_weight: torch.Tensor, + kv_norm_eps: float, + rotary_cos_sin: torch.Tensor, + sequence_length: torch.Tensor, + num_contexts: int, + seq_len_q: int, + kv_cache_block_offsets: torch.Tensor, + host_kv_cache_pool_pointers: torch.Tensor, + host_kv_cache_pool_mapping: torch.Tensor, + layer_idx: int, + tokens_per_block: int, + attention_window_size: int, + mla_bmm1_scale: torch.Tensor, + mla_bmm2_scale: torch.Tensor, + softmax_scale: float, + trigger_at_entry: bool = False, + ) -> None: + """The lengths the op derives positions from without checking. + + Only ``latent``'s row count, ``sequence_length``, ``num_contexts``, ``seq_len_q``, + ``rotary_cos_sin``'s size, ``kv_cache_block_offsets``' width and ``tokens_per_block`` + are read. + """ + num_generations = latent.shape[0] // seq_len_q + if num_generations == 0: + return + lens = sequence_length[num_contexts : num_contexts + num_generations].long() + # A sequence shorter than its step leaves rows at negative positions, which the kernel + # skips: those tokens never reach the cache. + assert bool((lens >= seq_len_q).all()), ( + "sequence_length must count this step's seq_len_q tokens; a shorter sequence " + "leaves its rows unwritten" + ) + last_position = int(lens.max()) - 1 + # The table is read at the position without a bound. + num_positions = rotary_cos_sin.numel() // (2 * ROPE_DIM) + assert last_position < num_positions, ( + f"rotary_cos_sin covers {num_positions} positions; position {last_position} " + "reads past it" + ) + # A block past the offsets' width reads another sequence's offsets and writes its page. + width = kv_cache_block_offsets.shape[-1] + assert last_position // tokens_per_block < width, ( + f"kv_cache_block_offsets has {width} blocks per sequence; position {last_position} " + "needs more" + ) + + +mla_kv_half_append = GlmMlaKvHalfAppend() diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_mla_kv_half_append.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_mla_kv_half_append.py new file mode 100644 index 000000000000..07147b5ef0d2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_mla_kv_half_append.py @@ -0,0 +1,290 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM-5.2 tp16 mla_kv_half_append entry. + +The paged state is real: a DSACacheManagerV2 over GLM-5.2's latent rows (FP8, 576 per token, +64-token pages) with several layers, the attention metadata of one generation step after a +context, and the MLA rope table at GLM-5.2's rope_theta. The rows the op appends are read back +where main's `convert_req_index_to_global` places their positions -- the rows main's built-in +sparse MLA decode reads -- and gated against the entry's reference; no other byte of the pool +may change. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.attention.mla_kv_half_append import ( # noqa: E501 + KV_LORA_RANK, + LATENT_OFFSET, + ROPE_DIM, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.attention.mla_kv_half_append import ( # noqa: E501 + mla_kv_half_append as op, +) +from tensorrt_llm._torch.attention.backends.interface import RopeParams +from tensorrt_llm._torch.attention.backends.sparse.dsa.cache_manager import DSACacheManagerV2 +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig, KvCacheConfig +from tensorrt_llm.mapping import Mapping + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or get_sm_version() != 107, + reason="the GLM-5.2 tp16 target's ops are certified on sm_107", +) + +HEAD_DIM = KV_LORA_RANK + ROPE_DIM +MAX_SEQ_LEN = 8192 +#: One E4M3 step: 2**-3 relative at the bottom of a binade, 2**-9 among the subnormals. +_E4M3_STEP_RTOL = 2**-3 +_E4M3_SUBNORMAL_STEP = 2**-9 +#: How much of an appended row may miss bit-exactness. See `_compare`. +_MAX_INEXACT_FRACTION = 1e-3 + + +def _compare(got: torch.Tensor, want: torch.Tensor) -> None: + """One E4M3 step, and at most 0.1 % of the bytes off bit-exactness. + + Kernel and reference evaluate the same fp32 expressions and round once; they differ in the + order of the 512-term sum of squares, in rsqrtf against torch.rsqrt (2 ulp), and in FMA + contraction of the rotation -- about 1e-7 relative, which crosses an E4M3 rounding boundary + (3 mantissa bits) for about one value in a million, and then by one step. A reference that + rounds the normalized latent and the rotated K to BF16 first (mla_rope_generation's KV + half) misses bit-exactness on about 3 % of the bytes, so the floor tells the two apart; a + wrong weight, epsilon, pair layout or position moves most bytes by several steps. + """ + assert got.shape == want.shape, (got.shape, want.shape) + torch.testing.assert_close( + got.float(), want.float(), rtol=_E4M3_STEP_RTOL, atol=_E4M3_SUBNORMAL_STEP + ) + inexact = int((got.view(torch.uint8) != want.view(torch.uint8)).sum()) + allowed = max(1, int(_MAX_INEXACT_FRACTION * got.numel())) + assert inexact <= allowed, f"{inexact}/{got.numel()} bytes not bit-exact (allowed {allowed})" + + +def _fill_e4m3_(dst: torch.Tensor, gen: torch.Generator) -> None: + """Fill a contiguous uint8 view with E4M3 values of N(0, 1).""" + flat = dst.reshape(-1) + chunk = 1 << 26 + for start in range(0, flat.numel(), chunk): + stop = min(start + chunk, flat.numel()) + values = torch.randn(stop - start, generator=gen, device=flat.device) + flat[start:stop].copy_(values.to(torch.float8_e4m3fn).view(torch.uint8)) + + +class _DsaEnv: + """A DSACacheManagerV2 over GLM-5.2's latent rows and one generation step's metadata.""" + + def __init__(self, num_layers: int, tokens_per_block: int, rope_theta: float) -> None: + self.tokens_per_block = tokens_per_block + self.max_batch_size = 2 + self.manager = DSACacheManagerV2( + kv_cache_config=KvCacheConfig( + max_tokens=self.max_batch_size * (MAX_SEQ_LEN + tokens_per_block), + enable_block_reuse=False, + ), + kv_cache_type=CacheType.SELFKONLY, + num_layers=num_layers, + num_kv_heads=1, + head_dim=HEAD_DIM, + tokens_per_block=tokens_per_block, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=self.max_batch_size, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=DataType.FP8, + sparse_attention_config=DeepSeekSparseAttentionConfig( + index_head_dim=128, index_n_heads=32, index_topk=2048 + ), + ) + rope = RopeParams( + dim=ROPE_DIM, theta=rope_theta, max_positions=MAX_SEQ_LEN, duplicate_data=True + ) + _, self.rotary_cos_sin = rope.create_rope_const_params() + + def generation_step(self, cached: list, seq_len_q: int) -> TrtllmAttentionMetadata: + """Requests 0 .. n-1 with `cached` tokens each, grown by one step of `seq_len_q`.""" + request_ids = list(range(len(cached))) + self.manager.add_dummy_requests(request_ids, token_nums=cached) + for rid in request_ids: + cache = self.manager.kv_cache_map[rid] + assert cache.resize(cache.capacity + seq_len_q) + metadata = TrtllmAttentionMetadata( + max_num_requests=self.max_batch_size, + max_num_tokens=8192, + kv_cache_manager=self.manager, + ) + metadata.seq_lens = torch.tensor([seq_len_q] * len(cached), dtype=torch.int) + metadata.num_contexts = 0 + metadata.request_ids = request_ids + metadata.prompt_lens = cached + metadata.kv_cache_params = KVCacheParams( + use_cache=True, num_cached_tokens_per_seq=cached + ) + metadata.prepare() + return metadata + + def pool_rows(self) -> torch.Tensor: + """The pool as [rows, 576] bytes, the rows convert_req_index_to_global counts.""" + pool = self.manager.get_unique_primary_pool() + return pool.view(torch.uint8).reshape(-1, HEAD_DIM) + + def global_rows(self, kv_lens: list, positions: torch.Tensor, seq_len_q: int, layer: int): + """Main's pool rows of each generation token's position in `layer`.""" + tpb = self.tokens_per_block + slots = self.manager.get_pool_block_indices( + len(kv_lens), request_ids=list(range(len(kv_lens))), num_contexts=0 + ) + blocks = torch.tensor([-(-n // tpb) for n in kv_lens]) + width = int(blocks.max()) + table = slots[:, :width].to(torch.int32).clone() + table[torch.arange(width)[None, :] >= blocks[:, None]] = -1 + page_index_scale, layer_offset = self.manager.get_primary_pool_page_index_params(layer) + req_idx = (torch.arange(positions.shape[0], device="cuda") // seq_len_q).to(torch.int32) + rows = torch.ops.trtllm.convert_req_index_to_global( + req_idx, + table.cuda(), + positions.to(torch.int32)[:, None].contiguous(), + tpb, + 1, + page_index_scale * tpb, + layer_offset, + ) + return rows[:, 0].long() + + def shutdown(self) -> None: + self.manager.shutdown() + + +def _step_args(env: _DsaEnv, spec: dict, gen: torch.Generator) -> tuple: + """One generation step's call: the metadata, a fused-A output and the scale buffers.""" + cached = list(spec["cached"]) + seq_len_q = spec["seq_len_q"] + metadata = env.generation_step(cached, seq_len_q) + num_tokens = len(cached) * seq_len_q + fused_a = torch.randn(num_tokens, spec["latent_width"], generator=gen, device="cuda").to( + torch.bfloat16 + ) + kv_norm_weight = (1.0 + 0.1 * torch.randn(KV_LORA_RANK, generator=gen, device="cuda")).to( + torch.bfloat16 + ) + return ( + fused_a[:, LATENT_OFFSET : LATENT_OFFSET + HEAD_DIM], + kv_norm_weight, + spec["kv_norm_eps"], + env.rotary_cos_sin, + metadata.kv_lens_cuda_runtime, + 0, + seq_len_q, + metadata.kv_cache_block_offsets, + env.manager.kv_cache_pool_pointers, + env.manager.kv_cache_pool_mapping, + spec["layer_idx"], + spec["tokens_per_block"], + env.manager.max_seq_len, + torch.full((2,), float("nan"), dtype=torch.float32, device="cuda"), + torch.full((1,), float("nan"), dtype=torch.float32, device="cuda"), + spec["softmax_scale"], + spec["trigger_at_entry"], + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by name + certifying(op) + spec = cell.spec + seed = zlib.crc32(cell.why.encode()) % 2**31 + gen = torch.Generator(device="cuda").manual_seed(seed) + env = _DsaEnv(spec["num_layers"], spec["tokens_per_block"], spec["rope_theta"]) + try: + args = _step_args(env, spec, gen) + latent, bmm1, bmm2 = args[0], args[13], args[14] + pool_rows = env.pool_rows() + _fill_e4m3_(pool_rows, gen) + pool_before = pool_rows.clone() + latent_before = latent.clone() + + op(*args) + torch.cuda.synchronize() + want_rows, positions, want_bmm1, want_bmm2 = op.reference(*args) + + # The position of each row: this step's tokens are the last seq_len_q of each sequence. + seq_len_q = spec["seq_len_q"] + kv_lens = [c + seq_len_q for c in spec["cached"]] + assert positions.tolist() == [n - seq_len_q + t for n in kv_lens for t in range(seq_len_q)] + + # Each row lands where main's layout puts its position in this layer. + rows = env.global_rows(kv_lens, positions, seq_len_q, spec["layer_idx"]) + assert rows.unique().numel() == rows.numel() and bool((rows >= 0).all()) + _compare(pool_rows[rows].view(torch.float8_e4m3fn), want_rows) + + # And nothing else moves: not another layer, slot or position, not the latent. + changed = (pool_rows != pool_before).any(dim=-1).nonzero().flatten() + assert set(changed.tolist()) <= set(rows.tolist()), "the op wrote outside its rows" + assert torch.equal(latent, latent_before) + + # The attention scales, bit for bit: one fp32 product. + assert torch.equal(bmm1, want_bmm1), (bmm1, want_bmm1) + assert torch.equal(bmm2, want_bmm2), (bmm2, want_bmm2) + + # Rerunning the prepared step rewrites the same bytes. + appended = pool_rows[rows].clone() + op(*args) + torch.cuda.synchronize() + assert torch.equal(pool_rows[rows], appended) + finally: + env.shutdown() + + +def _guard_args(**overrides) -> tuple: + """The first cell's call on a fresh environment, with arguments replaced by position.""" + spec = op.CELLS[0].spec + gen = torch.Generator(device="cuda").manual_seed(0) + env = _DsaEnv(spec["num_layers"], spec["tokens_per_block"], spec["rope_theta"]) + args = list(_step_args(env, spec, gen)) + names = ("sequence_length", "rotary_cos_sin", "kv_cache_block_offsets") + positions = {"sequence_length": 4, "rotary_cos_sin": 3, "kv_cache_block_offsets": 7} + for name, value in overrides.items(): + assert name in names + args[positions[name]] = value(args[positions[name]]) + return env, tuple(args) + + +def test_guard_refuses_a_sequence_shorter_than_its_step() -> None: + """Rows at negative positions are skipped: those tokens never reach the cache.""" + env, args = _guard_args(sequence_length=lambda t: torch.full_like(t, 5)) + try: + with pytest.raises(AssertionError, match="unwritten"): + with validating(op): + op(*args) + finally: + env.shutdown() + + +def test_guard_refuses_a_rope_table_shorter_than_the_positions() -> None: + """The table is read at the position without a bound.""" + env, args = _guard_args(rotary_cos_sin=lambda t: t.reshape(-1)[: 100 * 2 * ROPE_DIM]) + try: + with pytest.raises(AssertionError, match="reads past it"): + with validating(op): + op(*args) + finally: + env.shutdown() + + +def test_guard_refuses_block_offsets_narrower_than_the_positions() -> None: + """A block past the offsets' width is read from the next sequence's offsets.""" + env, args = _guard_args(kv_cache_block_offsets=lambda t: t[..., :8].contiguous()) + try: + with pytest.raises(AssertionError, match="needs more"): + with validating(op): + op(*args) + finally: + env.shutdown() From da28d087aaa13b532b662a1be0e772abe0cd7ccc Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:08:13 -0700 Subject: [PATCH 20/53] Add the GLM-5.2 tp16 target's sparse MLA flash-decode op glm_sparse_mla_flash_decode (kernels/glmMoeDsa/glmSparseMlaFlashDecode.cu, thop/glmSparseMlaFlashDecodeOp.cpp) runs the DSA sparse MLA generation attention of 4 query heads per rank as flash decoding over the gathered latent rows of a top-k of up to 2048: one cluster of 16 CTAs per query token, FP16 mma.sync on exactly unpacked E4M3 operands, P in E4M3 as the TRTLLM-Gen generation kernels feed their second GEMM, partial outputs pushed into the merging CTA during P.V, and the v_b BMM folded into the [tokens, 4 * 256] output. With a block table the kernel maps top-k positions to pool rows itself, with convert_req_index_to_global's arithmetic; with gather_before_wait it also reads the indices, the table, the KV lengths and the rows of earlier steps before its grid-dependency wait, for layers that share an earlier layer's top-k. The modeling_v2 entry ops/attention/sparse_mla_flash_decode.py states the op, its reference (with the E4M3 rounding of P per 128-slot split) and its cells: the verify full-indexer and shared-indexer layers, MTP draft step 0 and steps 1-6, a short context and two sequences. Its sm_107 test also replays the op from a captured CUDA graph and runs both GLM attention ops on a DSACacheManagerV2 pool against convert_req_index_to_global. Signed-off-by: Dongfeng Yu --- .../glmMoeDsa/glmSparseMlaFlashDecode.cu | 818 ++++++++++++++++++ .../glmMoeDsa/glmSparseMlaFlashDecode.h | 105 +++ cpp/tensorrt_llm/thop/CMakeLists.txt | 3 +- .../thop/glmSparseMlaFlashDecodeOp.cpp | 177 ++++ .../ops/attention/sparse_mla_flash_decode.py | 363 ++++++++ ...est_modeling_v2_sparse_mla_flash_decode.py | 563 ++++++++++++ 6 files changed, 2028 insertions(+), 1 deletion(-) create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.cu create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.h create mode 100644 cpp/tensorrt_llm/thop/glmSparseMlaFlashDecodeOp.cpp create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/sparse_mla_flash_decode.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_sparse_mla_flash_decode.py diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.cu new file mode 100644 index 000000000000..795a11fcbce2 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.cu @@ -0,0 +1,818 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ +namespace +{ + +// The MMAs run 8 head columns: the kHeads query heads in columns 0 .. kHeads - 1, zero Q rows in the others. +constexpr int kMmaHeads = 8; +constexpr int kHeads = kGlmSparseMlaNumHeads; +constexpr int kHeadDimQk = kGlmSparseMlaHeadDimQk; +constexpr int kHeadDimV = kGlmSparseMlaHeadDimV; +// One cluster of kSplits CTAs per query token; CTA s owns indices [s * kTokensPerCta, (s + 1) * kTokensPerCta). +constexpr int kSplits = 16; +constexpr int kTokensPerCta = kGlmSparseMlaMaxTopK / kSplits; +constexpr int kThreads = 256; +constexpr int kWarps = kThreads / 32; +constexpr int kTokensPerWarp = kTokensPerCta / kWarps; +static_assert(kTokensPerWarp == 16, "a warp gathers one 16-token MMA tile, one token row per lane"); +static_assert(kHeads < kMmaHeads, "the query heads fill part of the MMA head columns"); +// A token row: 576 E4M3 values. In shared memory the rows are 592 bytes apart, so the eight 16-byte rows of every +// ldmatrix phase fall into distinct banks. +constexpr int kRowBytes = kHeadDimQk; +constexpr int kSmemRowBytes = kRowBytes + 16; +// P (E4M3) per MMA head column, padded the same way. +constexpr int kPRowBytes = kTokensPerCta + 16; +// Output columns merged by each CTA of the cluster. +constexpr int kDvPerCta = kHeadDimV / kSplits; +static_assert(kDvPerCta % 16 == 0, "each output slice holds whole 16-column MMA blocks"); +// P is scaled to the E4M3 range before the second MMA; the normalization divides it out. +constexpr float kPScale = 448.f; + +// v_b BMM: CTA r of the cluster owns v in [r * kVPerCta, (r + 1) * kVPerCta) of every head. +constexpr int kVHeadDim = kGlmSparseMlaVHeadDim; +constexpr int kVPerCta = kVHeadDim / kSplits; +static_assert(kVPerCta == 16, "a warp computes one head's 16 rows of the v_b slice as one m16n8k16 row tile"); +// The v_b_proj rows [h][r * 16, r * 16 + 16)[0, 512) of one head: 16 KiB, contiguous in global memory. +constexpr int kVbHeadBytes = kVPerCta * kHeadDimV * 2; + +// Shared memory: token rows, Q rows, P, the per-warp max / sum, the receive buffers of the cluster merge, the v_b +// slice, the CTA's merged output slice, the row's gathered BF16 output and the mbarriers. +constexpr int kOffQ = kTokensPerCta * kSmemRowBytes; +constexpr int kOffP = kOffQ + kMmaHeads * kSmemRowBytes; +constexpr int kOffRed = kOffP + kMmaHeads * kPRowBytes; +constexpr int kRedFloats = 2 * kWarps * kMmaHeads; +constexpr int kOffRecvO = kOffRed + kRedFloats * 4; +constexpr int kRecvOFloats = kSplits * kMmaHeads * kDvPerCta; +constexpr int kOffRecvStats = kOffRecvO + kRecvOFloats * 4; +constexpr int kRecvStatsFloats = kSplits * kMmaHeads * 2; +constexpr int kRecvBytes = (kRecvOFloats + kRecvStatsFloats) * 4; +constexpr int kOffVb = kOffRecvStats + kRecvStatsFloats * 4; +constexpr int kVbBytes = kHeads * kVbHeadBytes; +// The CTA's merged output slice in BF16, [8 head columns][32 columns]. +constexpr int kOffSlice = kOffVb + kVbBytes; +constexpr int kSliceBytes = kMmaHeads * kDvPerCta * 2; +// The row's merged output in BF16, [8 head columns][512]. +constexpr int kOffGatheredO = kOffSlice + kSliceBytes; +constexpr int kGatheredOBytes = kMmaHeads * kHeadDimV * 2; +constexpr int kOffBars = kOffGatheredO + kGatheredOBytes; +// One gather barrier per warp, the Q barrier, the receive barrier, the v_b slice barrier and the gathered-output +// barrier. +constexpr int kNumBars = kWarps + 4; +constexpr int kSmemBytes = kOffBars + kNumBars * 8; + +static_assert(kOffQ % 16 == 0 && kOffP % 16 == 0 && kOffRecvO % 16 == 0 && kOffVb % 16 == 0 && kOffSlice % 16 == 0 + && kOffGatheredO % 16 == 0 && kOffBars % 8 == 0, + "alignment"); +static_assert(kSplits * kMmaHeads <= kThreads, "one max / sum transfer per thread"); +static_assert(kMmaHeads * kDvPerCta == kThreads, "the merge maps one output element to each thread"); +static_assert(kDvPerCta == 32 && kSplits <= 32, "the merge maps warp h to head h of the slice, lane r to partial r"); +static_assert(kSmemBytes <= 227 * 1024, "shared memory per CTA"); + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + +__device__ __forceinline__ uint32_t smemAddress(void const* ptr) +{ + return static_cast(__cvta_generic_to_shared(ptr)); +} + +__device__ __forceinline__ void mbarrierInit(uint32_t bar, uint32_t count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(bar), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrierArriveExpectTx(uint32_t bar, uint32_t bytes) +{ + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(bar), "r"(bytes) : "memory"); +} + +__device__ __forceinline__ void mbarrierWait(uint32_t bar, uint32_t parity) +{ + asm volatile( + "{\n" + ".reg .pred complete;\n" + "GLM_SPARSE_MLA_FD_WAIT_%=:\n" + "mbarrier.try_wait.parity.shared::cta.b64 complete, [%0], %1;\n" + "@!complete bra GLM_SPARSE_MLA_FD_WAIT_%=;\n" + "}\n" ::"r"(bar), + "r"(parity) + : "memory"); +} + +// Wait with cluster-scope acquire: the receive barriers complete on bytes that other CTAs of the cluster wrote. +__device__ __forceinline__ void mbarrierWaitCluster(uint32_t bar, uint32_t parity) +{ + asm volatile( + "{\n" + ".reg .pred complete;\n" + "GLM_SPARSE_MLA_FD_WAIT_CLUSTER_%=:\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 complete, [%0], %1;\n" + "@!complete bra GLM_SPARSE_MLA_FD_WAIT_CLUSTER_%=;\n" + "}\n" ::"r"(bar), + "r"(parity) + : "memory"); +} + +__device__ __forceinline__ void bulkCopyGlobalToShared(uint32_t dst, void const* src, uint32_t bytes, uint32_t bar) +{ + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n" ::"r"(dst), + "l"(src), "r"(bytes), "r"(bar) + : "memory"); +} + +// Two floats into the shared memory of a CTA of the cluster; the bytes complete on that CTA's barrier. +__device__ __forceinline__ void storeRemote(uint32_t remoteAddr, float a, float b, uint32_t remoteBar) +{ + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.f32 [%0], {%1, %2}, [%3];\n" ::"r"(remoteAddr), + "f"(a), "f"(b), "r"(remoteBar) + : "memory"); +} + +// Sixteen bytes into the shared memory of a CTA of the cluster; the bytes complete on that CTA's barrier. +__device__ __forceinline__ void storeRemote16(uint32_t remoteAddr, uint4 value, uint32_t remoteBar) +{ + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];\n" ::"r"( + remoteAddr), + "r"(value.x), "r"(value.y), "r"(value.z), "r"(value.w), "r"(remoteBar) + : "memory"); +} + +__device__ __forceinline__ uint32_t mapToRank(uint32_t addr, uint32_t rank) +{ + uint32_t mapped; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;\n" : "=r"(mapped) : "r"(addr), "r"(rank)); + return mapped; +} + +__device__ __forceinline__ void clusterArriveRelaxed() +{ + asm volatile("barrier.cluster.arrive.relaxed.aligned;\n" ::: "memory"); +} + +__device__ __forceinline__ void clusterWait() +{ + asm volatile("barrier.cluster.wait.acquire.aligned;\n" ::: "memory"); +} + +__device__ __forceinline__ void ldmatrixX4(uint32_t addr, uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) + : "r"(addr)); +} + +__device__ __forceinline__ void ldmatrixX4Trans(uint32_t addr, uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) + : "r"(addr)); +} + +__device__ __forceinline__ void ldmatrixX2(uint32_t addr, uint32_t& r0, uint32_t& r1) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" : "=r"(r0), "=r"(r1) : "r"(addr)); +} + +// D += A B with A 16 x 16 (row major) and B 16 x 8 (column major), BF16 inputs and fp32 accumulation. +__device__ __forceinline__ void mmaBf16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +// D += A B with A 16 x 16 (row major) and B 16 x 8 (column major), FP16 inputs and fp32 accumulation. +__device__ __forceinline__ void mmaF16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +// The two E4M3 values in the low 16 bits of `pair` as two FP16 values (exact). +__device__ __forceinline__ uint32_t e4m3x2ToHalf2(uint32_t pair) +{ + __half2_raw const h = __nv_cvt_fp8x2_to_halfraw2(static_cast<__nv_fp8x2_storage_t>(pair & 0xffffU), __NV_E4M3); + return static_cast(h.x) | (static_cast(h.y) << 16); +} + +// The B fragment of an E4M3 m16n8k32 MMA ({b0, b1}) unpacked for the two FP16 MMAs of mmaE4m3AsF16: low halves, then +// high halves. +__device__ __forceinline__ void unpackE4m3B(uint32_t (&bh)[4], uint32_t b0, uint32_t b1) +{ + bh[0] = e4m3x2ToHalf2(b0); + bh[1] = e4m3x2ToHalf2(b1); + bh[2] = e4m3x2ToHalf2(b0 >> 16); + bh[3] = e4m3x2ToHalf2(b1 >> 16); +} + +// D += A B for an E4M3 m16n8k32 MMA with A in its E4M3 fragment registers and B unpacked beforehand (unpackE4m3B): +// an FP16 m16n8k16 MMA on the low 16-bit halves of the A registers, then one on the high halves, chained on D (the +// SM100 family has no E4M3 mma.sync; this is the sequence the compiler lowers one to). +__device__ __forceinline__ void mmaE4m3AsF16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t const (&bh)[4]) +{ + mmaF16(d, e4m3x2ToHalf2(a0), e4m3x2ToHalf2(a1), e4m3x2ToHalf2(a2), e4m3x2ToHalf2(a3), bh[0], bh[1]); + mmaF16(d, e4m3x2ToHalf2(a0 >> 16), e4m3x2ToHalf2(a1 >> 16), e4m3x2ToHalf2(a2 >> 16), e4m3x2ToHalf2(a3 >> 16), + bh[2], bh[3]); +} + +__device__ __forceinline__ uint8_t toE4m3(float x) +{ + return static_cast(__nv_cvt_float_to_fp8(x, __NV_SATFINITE, __NV_E4M3)); +} + +// Position of the CTA's token t in the k order of the P.V MMA: within every 16 tokens, k slots 4c .. 4c + 3 hold the +// tokens {2c, 2c + 1, 2c + 8, 2c + 9}, the order in which ldmatrix.trans of the token-major V rows plus two byte +// permutes deliver a V^T fragment. +__device__ __forceinline__ int permutedSlot(int t) +{ + int const u = t & 15; + int const slot16 = (u < 8) ? ((u >> 1) * 4 + (u & 1)) : (((u - 8) >> 1) * 4 + 2 + (u & 1)); + return (t & ~15) + slot16; +} + +// The row of one lane of a warp's gather, found before the grid-dependency wait (kGatherBeforeWait). +struct EarlyGatherRow +{ + int64_t row; + bool valid; + // A position of this step's tokens: its row is written by the KV append this grid depends on. + bool fresh; +}; + +// Lane l < 16 of the warp maps index slot `slot` of query token `tokenQ` (a position) to its pool row with the remap +// arithmetic, and bulk-copies the row unless it holds one of this step's tokens. Masked lanes copy row 0, as after the +// wait. Every read here (indices, block table, KV length, rows of earlier steps) precedes this grid's launch. +__device__ __forceinline__ EarlyGatherRow gatherRowBeforeWait(GlmSparseMlaFlashDecodeParams const& params, int tokenQ, + int slot, int lane, uint32_t dst, uint32_t barGather) +{ + int position = -1; + if (lane < kTokensPerWarp && slot < params.topK) + { + position = __ldg(params.indices + static_cast(tokenQ) * params.indicesStride + slot); + } + int const seqIdx = tokenQ / params.seqLenQ; + int const tokenInSeq = tokenQ - seqIdx * params.seqLenQ; + int const block = position / params.remapTokensPerBlock; + int page = -1; + if (position >= 0 && block < params.remapMaxBlocks) + { + page = __ldg(params.remapBlockTable + seqIdx * params.remapBlockTableStride + block); + } + int const inBlock = position - block * params.remapTokensPerBlock; + int const rowRaw = page >= 0 ? page * params.remapStrideFactor + inBlock + params.remapLayerOffset : -1; + int const seqLenKv = __ldg(params.seqLensKv + seqIdx); + int const numValid = min(seqLenKv - (params.seqLenQ - 1 - tokenInSeq), params.topK); + EarlyGatherRow out; + out.valid = lane < kTokensPerWarp && slot < numValid && rowRaw >= 0; + out.row = out.valid ? rowRaw : 0; + out.fresh = out.valid && position >= seqLenKv - params.seqLenQ; + if (lane == 0) + { + mbarrierArriveExpectTx(barGather, kTokensPerWarp * kRowBytes); + } + __syncwarp(); + if (lane < kTokensPerWarp && !out.fresh) + { + bulkCopyGlobalToShared(dst, static_cast(params.kvPool) + out.row * kRowBytes, kRowBytes, + barGather); + } + return out; +} + +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + +// kRemap: the indices are positions, mapped to pool rows through the block table. kGatherBeforeWait (with kRemap): the +// indices, the block table and the KV length are read, and the rows of the positions before this step's tokens +// gathered, before the grid-dependency wait; the rows of this step's tokens after it. +template +__global__ void __launch_bounds__(kThreads, 1) glmSparseMlaFlashDecodeKernel(GlmSparseMlaFlashDecodeParams const params) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + static_assert(kRemap || !kGatherBeforeWait, "the early gather tells this step's tokens by their positions"); + extern __shared__ __align__(1024) uint8_t smem[]; + uint8_t* sKv = smem; + uint8_t* sQ = smem + kOffQ; + uint8_t* sP = smem + kOffP; + float* sRedMax = reinterpret_cast(smem + kOffRed); + float* sRedSum = sRedMax + kWarps * kMmaHeads; + float* recvO = reinterpret_cast(smem + kOffRecvO); + float* recvStats = reinterpret_cast(smem + kOffRecvStats); + uint8_t* sSlice = smem + kOffSlice; + uint8_t* sGatheredO = smem + kOffGatheredO; + uint64_t* bars = reinterpret_cast(smem + kOffBars); + + int const tid = threadIdx.x; + int const warp = tid / 32; + int const lane = tid % 32; + // The cluster spans gridDim.x, so the split index is the CTA's rank in the cluster. + int const split = blockIdx.x; + int const tokenQ = blockIdx.y; + + uint32_t const barGather = smemAddress(bars + warp); + uint32_t const barQ = smemAddress(bars + kWarps); + uint32_t const barRecv = smemAddress(bars + kWarps + 1); + uint32_t const barVb = smemAddress(bars + kWarps + 2); + uint32_t const barGatheredO = smemAddress(bars + kWarps + 3); + { + // Q rows kHeads .. 7 feed the MMAs' unused head columns: zero, before the barrier below orders them for the QK. + constexpr int kZeroChunks = (kMmaHeads - kHeads) * kRowBytes / 16; + static_assert(kZeroChunks <= kThreads && kRowBytes % 16 == 0, "one 16-byte zero store per thread"); + if (tid < kZeroChunks) + { + int const r = kHeads + tid / (kRowBytes / 16); + int const chunk = tid % (kRowBytes / 16); + *reinterpret_cast(sQ + r * kSmemRowBytes + 16 * chunk) = make_uint4(0, 0, 0, 0); + } + } + if (tid == 0) + { + for (int i = 0; i < kNumBars; ++i) + { + mbarrierInit(smemAddress(bars + i), 1); + } + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + mbarrierArriveExpectTx(barRecv, kRecvBytes); + mbarrierArriveExpectTx(barGatheredO, kGatheredOBytes); + } + __syncthreads(); + // Paired with the wait before the first remote store: no CTA writes into a receive buffer whose barrier other + // CTAs of the cluster have not initialized yet. + clusterArriveRelaxed(); + + [[maybe_unused]] EarlyGatherRow early{}; + if constexpr (kGatherBeforeWait) + { + if (warp == 0) + { + // The block-table row of the token's sequence into L2 while the indices load: the remap below reads it + // next. The CTAs of a cluster prefetch interleaved 128-byte lines of the row. + char const* tableRow = reinterpret_cast( + params.remapBlockTable + (tokenQ / params.seqLenQ) * params.remapBlockTableStride); + int64_t const rowBytes = static_cast(params.remapMaxBlocks) * sizeof(int32_t); + for (int64_t off = (static_cast(split) * 32 + lane) * 128; off < rowBytes; + off += kSplits * 32 * 128) + { + asm volatile("prefetch.global.L2 [%0];\n" ::"l"(tableRow + off)); + } + } + int const slotEarly = split * kTokensPerCta + warp * kTokensPerWarp + lane; + early = gatherRowBeforeWait(params, tokenQ, slotEarly, lane, + smemAddress(sKv + (warp * kTokensPerWarp + lane) * kSmemRowBytes), barGather); + } + + // Q, the indices, the KV pool rows and the scales come from earlier kernels. + cudaGridDependencySynchronize(); + + if (warp == 0) + { + if (lane == 0) + { + mbarrierArriveExpectTx(barQ, kHeads * kRowBytes); + } + __syncwarp(); + if (lane < kHeads) + { + bulkCopyGlobalToShared(smemAddress(sQ + lane * kSmemRowBytes), + static_cast(params.q) + (static_cast(tokenQ) * kHeads + lane) * kRowBytes, + kRowBytes, barQ); + } + } + + // Lane l of warp w gathers token tokBase + l of the CTA. The index is read before the KV length is known, so the + // two loads share one round trip. + int const tokBase = warp * kTokensPerWarp; + uint32_t validMask; + if constexpr (kGatherBeforeWait) + { + // The rows of this step's tokens, which the KV append wrote. + if (lane < kTokensPerWarp && early.fresh) + { + bulkCopyGlobalToShared(smemAddress(sKv + (tokBase + lane) * kSmemRowBytes), + static_cast(params.kvPool) + early.row * kRowBytes, kRowBytes, barGather); + } + validMask = __ballot_sync(0xffffffffU, early.valid); + } + else + { + int const slot = split * kTokensPerCta + tokBase + lane; + int indexRaw = -1; + if (lane < kTokensPerWarp && slot < params.topK) + { + indexRaw = __ldg(params.indices + static_cast(tokenQ) * params.indicesStride + slot); + } + int const seqIdx = tokenQ / params.seqLenQ; + int const tokenInSeq = tokenQ - seqIdx * params.seqLenQ; + if constexpr (kRemap) + { + // The index is a position in the sequence: map it to its pool row through the block table. + int const block = indexRaw / params.remapTokensPerBlock; + int page = -1; + if (indexRaw >= 0 && block < params.remapMaxBlocks) + { + page = __ldg(params.remapBlockTable + seqIdx * params.remapBlockTableStride + block); + } + int const inBlock = indexRaw - block * params.remapTokensPerBlock; + indexRaw = page >= 0 ? page * params.remapStrideFactor + inBlock + params.remapLayerOffset : -1; + } + int const numValid = min(__ldg(params.seqLensKv + seqIdx) - (params.seqLenQ - 1 - tokenInSeq), params.topK); + bool const valid = lane < kTokensPerWarp && slot < numValid && indexRaw >= 0; + // Masked lanes still load a row (row 0 of the pool) so that every gather barrier expects the same bytes. + int64_t const row = valid ? indexRaw : 0; + if (lane == 0) + { + mbarrierArriveExpectTx(barGather, kTokensPerWarp * kRowBytes); + } + __syncwarp(); + if (lane < kTokensPerWarp) + { + bulkCopyGlobalToShared(smemAddress(sKv + (tokBase + lane) * kSmemRowBytes), + static_cast(params.kvPool) + row * kRowBytes, kRowBytes, barGather); + } + validMask = __ballot_sync(0xffffffffU, valid); + } + float const scaleSoftmaxLog2 = __ldg(params.scaleSoftmaxLog2); + float const outputScale = __ldg(params.outputScale); + + // Fragment coordinates of mma.sync: g = row group, c = column pair. + int const g = lane / 4; + int const c = lane % 4; + + // S^T = K Q^T for the warp's 16 tokens and the 8 head columns; two accumulators over alternating k steps. The + // query fragments of every k step are unpacked once, while the warp's token rows are in flight. + mbarrierWait(barQ, 0); + uint32_t qh[kHeadDimQk / 32][4]; + { + uint32_t const qAddr = smemAddress(sQ + (lane % 8) * kSmemRowBytes + 16 * ((lane / 8) % 2)); +#pragma unroll + for (int kd = 0; kd < kHeadDimQk / 32; ++kd) + { + uint32_t b0, b1; + ldmatrixX2(qAddr + 32 * kd, b0, b1); + unpackE4m3B(qh[kd], b0, b1); + } + } + mbarrierWait(barGather, 0); + float s[4]; + { + float acc[2][4] = {{0.f, 0.f, 0.f, 0.f}, {0.f, 0.f, 0.f, 0.f}}; + uint32_t const kAddr + = smemAddress(sKv + (tokBase + lane % 8 + 8 * ((lane / 8) % 2)) * kSmemRowBytes + 16 * (lane / 16)); +#pragma unroll + for (int kd = 0; kd < kHeadDimQk / 32; ++kd) + { + uint32_t a0, a1, a2, a3; + ldmatrixX4(kAddr + 32 * kd, a0, a1, a2, a3); + mmaE4m3AsF16(acc[kd % 2], a0, a1, a2, a3, qh[kd]); + } + bool const valid0 = (validMask >> g) & 1U; + bool const valid1 = (validMask >> (g + 8)) & 1U; + s[0] = valid0 ? (acc[0][0] + acc[1][0]) * scaleSoftmaxLog2 : -INFINITY; + s[1] = valid0 ? (acc[0][1] + acc[1][1]) * scaleSoftmaxLog2 : -INFINITY; + s[2] = valid1 ? (acc[0][2] + acc[1][2]) * scaleSoftmaxLog2 : -INFINITY; + s[3] = valid1 ? (acc[0][3] + acc[1][3]) * scaleSoftmaxLog2 : -INFINITY; + } + + // The CTA's max per head column; this thread holds columns 2c and 2c + 1. + float max0 = fmaxf(s[0], s[2]); + float max1 = fmaxf(s[1], s[3]); +#pragma unroll + for (int offset = 4; offset < 32; offset *= 2) + { + max0 = fmaxf(max0, __shfl_xor_sync(0xffffffffU, max0, offset)); + max1 = fmaxf(max1, __shfl_xor_sync(0xffffffffU, max1, offset)); + } + if (g == 0) + { + sRedMax[warp * kMmaHeads + 2 * c] = max0; + sRedMax[warp * kMmaHeads + 2 * c + 1] = max1; + } + // Every warp observed its own gather barrier before its QK; this barrier orders all token rows before the P.V + // reads of the other warps. + __syncthreads(); + if (warp == 1) + { + // v_b_proj[h][16 r, 16 r + 16)[0, 512) for every head h: a weight, so it may be read at any point. It is + // issued once every token row has landed and arrives during QK / P.V. + if (lane == 0) + { + mbarrierArriveExpectTx(barVb, kVbBytes); + } + __syncwarp(); + if (lane < kHeads) + { + bulkCopyGlobalToShared(smemAddress(smem + kOffVb + lane * kVbHeadBytes), + static_cast(params.vbProj) + + (static_cast(lane) * kVHeadDim + split * kVPerCta) * kHeadDimV * 2, + kVbHeadBytes, barVb); + } + } + float ctaMax0 = -INFINITY; + float ctaMax1 = -INFINITY; +#pragma unroll + for (int w = 0; w < kWarps; ++w) + { + ctaMax0 = fmaxf(ctaMax0, sRedMax[w * kMmaHeads + 2 * c]); + ctaMax1 = fmaxf(ctaMax1, sRedMax[w * kMmaHeads + 2 * c + 1]); + } + // A CTA without valid tokens keeps P = 0. + float const offset0 = ctaMax0 == -INFINITY ? 0.f : ctaMax0; + float const offset1 = ctaMax1 == -INFINITY ? 0.f : ctaMax1; + + // P in E4M3 at the permuted k slots, and the per-head sums of the fp32 P. + { + float const p00 = exp2f(s[0] - offset0); + float const p01 = exp2f(s[1] - offset1); + float const p10 = exp2f(s[2] - offset0); + float const p11 = exp2f(s[3] - offset1); + int const slot0 = permutedSlot(tokBase + g); + int const slot1 = permutedSlot(tokBase + g + 8); + sP[(2 * c) * kPRowBytes + slot0] = toE4m3(p00 * kPScale); + sP[(2 * c + 1) * kPRowBytes + slot0] = toE4m3(p01 * kPScale); + sP[(2 * c) * kPRowBytes + slot1] = toE4m3(p10 * kPScale); + sP[(2 * c + 1) * kPRowBytes + slot1] = toE4m3(p11 * kPScale); + float sum0 = p00 + p10; + float sum1 = p01 + p11; +#pragma unroll + for (int offset = 4; offset < 32; offset *= 2) + { + sum0 += __shfl_xor_sync(0xffffffffU, sum0, offset); + sum1 += __shfl_xor_sync(0xffffffffU, sum1, offset); + } + if (g == 0) + { + sRedSum[warp * kMmaHeads + 2 * c] = sum0; + sRedSum[warp * kMmaHeads + 2 * c + 1] = sum1; + } + } + __syncthreads(); + + // This CTA's max and sum per head column go to every CTA of the cluster now; the partial output follows during + // P.V. The merged slice and the gathered output have buffers of their own, so no merge can overwrite memory that a + // warp of the cluster still reads in P.V. + clusterWait(); + if (tid < kSplits * kMmaHeads) + { + int const dst = tid / kMmaHeads; + int const h = tid % kMmaHeads; + float ctaMax = -INFINITY; + float ctaSum = 0.f; +#pragma unroll + for (int w = 0; w < kWarps; ++w) + { + ctaMax = fmaxf(ctaMax, sRedMax[w * kMmaHeads + h]); + ctaSum += sRedSum[w * kMmaHeads + h]; + } + storeRemote(mapToRank(smemAddress(recvStats + (split * kMmaHeads + h) * 2), dst), ctaMax, ctaSum, + mapToRank(barRecv, dst)); + } + + // O^T = V^T P^T: warp w owns the 16-column output blocks 4w .. 4w + 3, one block at a time over all k steps, each + // block sent to its merging CTA while the next one is computed. ldmatrix.trans of 8 tokens x 16 bytes gives + // (token 2c, columns 2g and 2g + 1) and (token 2c + 1, same columns); the byte permutes regroup them into a V^T row + // (one column, four tokens), so fragment row g holds column 2g and row g + 8 holds column 2g + 1. + { + uint32_t const pAddr = smemAddress(sP + (lane % 8) * kPRowBytes + 16 * ((lane / 8) % 2)); + uint32_t const vAddr = smemAddress(sKv + lane * kSmemRowBytes + 16 * (4 * warp)); + uint32_t pb[kTokensPerCta / 32][4]; +#pragma unroll + for (int kStep = 0; kStep < kTokensPerCta / 32; ++kStep) + { + uint32_t b0, b1; + ldmatrixX2(pAddr + 32 * kStep, b0, b1); + unpackE4m3B(pb[kStep], b0, b1); + } +#pragma unroll + for (int bi = 0; bi < 4; ++bi) + { + float ob[4] = {0.f, 0.f, 0.f, 0.f}; +#pragma unroll + for (int kStep = 0; kStep < kTokensPerCta / 32; ++kStep) + { + uint32_t r0, r1, r2, r3; + ldmatrixX4Trans(vAddr + 32 * kStep * kSmemRowBytes + 16 * bi, r0, r1, r2, r3); + uint32_t const a0 = __byte_perm(r0, r1, 0x6420); + uint32_t const a1 = __byte_perm(r0, r1, 0x7531); + uint32_t const a2 = __byte_perm(r2, r3, 0x6420); + uint32_t const a3 = __byte_perm(r2, r3, 0x7531); + mmaE4m3AsF16(ob, a0, a1, a2, a3, pb[kStep]); + } + // Columns [32 r, 32 r + 32) of the partial output go to CTA r. + int const col = 16 * (4 * warp + bi) + 2 * g; + int const dst = col / kDvPerCta; + int const colInSlice = col % kDvPerCta; + uint32_t const remoteBar = mapToRank(barRecv, dst); + storeRemote(mapToRank(smemAddress(recvO + (split * kMmaHeads + 2 * c) * kDvPerCta + colInSlice), dst), + ob[0], ob[2], remoteBar); + storeRemote(mapToRank(smemAddress(recvO + (split * kMmaHeads + 2 * c + 1) * kDvPerCta + colInSlice), dst), + ob[1], ob[3], remoteBar); + } + } + // The dependents (the o_proj GEMM) wait for the whole grid before reading out; launching them now lets them stage + // their weights during the merge and the v_b product. + cudaTriggerProgrammaticLaunchCompletion(); + + // Merge the kSplits partials of this CTA's output slice: thread (h, column) of the 8 x 32 slice. Warp h merges head + // column h; lane r < kSplits computes the weight of partial r once, and the loop takes the weights from the lanes + // in partial order. + mbarrierWaitCluster(barRecv, 0); + int const h = tid / kDvPerCta; + int const colInSlice = tid % kDvPerCta; + float value; + { + float const laneMax = lane < kSplits ? recvStats[(lane * kMmaHeads + h) * 2] : -INFINITY; + float mergedMax = laneMax; +#pragma unroll + for (int offset = 16; offset > 0; offset /= 2) + { + mergedMax = fmaxf(mergedMax, __shfl_xor_sync(0xffffffffU, mergedMax, offset)); + } + float const laneAlpha = laneMax == -INFINITY ? 0.f : exp2f(laneMax - mergedMax); + float mergedSum = 0.f; + float acc = 0.f; +#pragma unroll + for (int r = 0; r < kSplits; ++r) + { + float const alpha = __shfl_sync(0xffffffffU, laneAlpha, r); + mergedSum += alpha * recvStats[(r * kMmaHeads + h) * 2 + 1]; + acc += alpha * recvO[(r * kMmaHeads + h) * kDvPerCta + colInSlice]; + } + value = mergedSum > 0.f ? acc / (mergedSum * kPScale) * outputScale : 0.f; + } + + // The merged slice in BF16 (the attention output the v_b product reads) goes to every CTA of the cluster. The merge + // maps thread tid to head column tid / 32, so each warp holds and sends one head row. + reinterpret_cast<__nv_bfloat16*>(sSlice)[h * kDvPerCta + colInSlice] = __float2bfloat16_rn(value); + __syncwarp(); + { + int const chunk = lane % 4; + uint4 const packed = *reinterpret_cast(sSlice + (h * kDvPerCta + 8 * chunk) * 2); + uint32_t const localAddr = smemAddress(sGatheredO + (h * kHeadDimV + split * kDvPerCta + 8 * chunk) * 2); +#pragma unroll + for (int k = 0; k < 2; ++k) + { + uint32_t const dst = lane / 4 + 8 * k; + storeRemote16(mapToRank(localAddr, dst), packed, mapToRank(barGatheredO, dst)); + } + } + mbarrierWaitCluster(barGatheredO, 0); + mbarrierWait(barVb, 0); + if (h >= kHeads) + { + // A head column without a v_b slice (warp-uniform: the merge maps warp h to head column h). + return; + } + + // out[t][h * 256 + 16 r + v] = sum_d O[h][d] v_b[h][16 r + v][d] for v = 0 .. 15: warp h, one m16n8k16 row tile of + // the 16 v rows (1 KiB apart). Per pair of k steps a thread loads 16 bytes of its rows g and g + 8; rows with odd g + // take the other 32-column half of every 64-column block, so the eight rows of a load cover all banks. B column n + // follows the same parity as rows with g = n, so the dot product of row g is D[g][g % 2] (thread c = 0). + { + int const parity = g % 2; + uint8_t const* rowG = smem + kOffVb + h * kVbHeadBytes + g * (kHeadDimV * 2) + 16 * c; + uint8_t const* rowG8 = rowG + 8 * (kHeadDimV * 2); + uint8_t const* oRow = sGatheredO + h * (kHeadDimV * 2) + 16 * c; + float acc[4][4]; +#pragma unroll + for (int a = 0; a < 4; ++a) + { +#pragma unroll + for (int i = 0; i < 4; ++i) + { + acc[a][i] = 0.f; + } + } +#pragma unroll + for (int kPair = 0; kPair < kHeadDimV / 32; ++kPair) + { + int const block = (kPair ^ parity) * 64; + uint4 const ra = *reinterpret_cast(rowG + block); + uint4 const rb = *reinterpret_cast(rowG8 + block); + uint4 const ob = *reinterpret_cast(oRow + block); + mmaBf16(acc[(2 * kPair) % 4], ra.x, rb.x, ra.y, rb.y, ob.x, ob.y); + mmaBf16(acc[(2 * kPair + 1) % 4], ra.z, rb.z, ra.w, rb.w, ob.z, ob.w); + } + if (c == 0) + { + float sums[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + sums[i] = ((acc[0][i] + acc[1][i]) + acc[2][i]) + acc[3][i]; + } + __nv_bfloat16* outRow = static_cast<__nv_bfloat16*>(params.out) + + static_cast(tokenQ) * params.outStride + h * kVHeadDim + split * kVPerCta; + outRow[g] = __float2bfloat16_rn(parity ? sums[1] : sums[0]); + outRow[g + 8] = __float2bfloat16_rn(parity ? sums[3] : sums[2]); + } + } +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +template +void launchGlmSparseMlaFlashDecode(GlmSparseMlaFlashDecodeParams const& params, cudaStream_t stream) +{ + auto const kernel = glmSparseMlaFlashDecodeKernel; + static bool const configured = [kernel]() + { + TLLM_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes)); + TLLM_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1)); + return true; + }(); + (void) configured; + + cudaLaunchConfig_t config; + config.gridDim = dim3(kSplits, params.numTokensQ, 1); + config.blockDim = dim3(kThreads, 1, 1); + config.dynamicSmemBytes = kSmemBytes; + config.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kSplits; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.attrs = attrs; + config.numAttrs = 2; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, params)); +} + +} // namespace + +void invokeGlmSparseMlaFlashDecode(GlmSparseMlaFlashDecodeParams const& params, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(params.topK > 0 && params.topK <= kGlmSparseMlaMaxTopK, + "GLM sparse MLA flash decode supports TopK up to %d, got %d", kGlmSparseMlaMaxTopK, params.topK); + TLLM_CHECK_WITH_INFO(params.seqLenQ > 0 && params.indicesStride >= params.topK && params.numTokensQ >= 0 + && params.numTokensQ % params.seqLenQ == 0, + "Invalid GLM sparse MLA geometry"); + TLLM_CHECK_WITH_INFO(params.vbProj != nullptr && params.out != nullptr && params.outStride >= kHeads * kVHeadDim, + "GLM sparse MLA flash decode needs v_b_proj and an output of %d columns per token", kHeads * kVHeadDim); + if (params.numTokensQ == 0) + { + return; + } + bool const remap = params.remapBlockTable != nullptr; + if (remap) + { + TLLM_CHECK_WITH_INFO(params.remapTokensPerBlock > 0 && params.remapMaxBlocks > 0 + && params.remapBlockTableStride >= params.remapMaxBlocks, + "Invalid GLM sparse MLA index remap geometry"); + } + TLLM_CHECK_WITH_INFO(!params.gatherBeforeWait || remap, + "The GLM sparse MLA flash decode gathers before its wait only with the index remap"); + if (remap && params.gatherBeforeWait) + { + launchGlmSparseMlaFlashDecode(params, stream); + } + else if (remap) + { + launchGlmSparseMlaFlashDecode(params, stream); + } + else + { + launchGlmSparseMlaFlashDecode(params, stream); + } +} + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.h new file mode 100644 index 000000000000..7620932c477f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.h @@ -0,0 +1,105 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +// The GLM-5.2 decode shape at TP16 the flash-decode sparse MLA serves: 4 query heads per rank sharing one latent KV +// head, QK head dim 576 (512 latent + 64 rope), latent V dim 512, a TopK of at most 2048 selected tokens per query +// token, and the v_b BMM (v_head_dim 256) folded into the output. +constexpr int32_t kGlmSparseMlaNumHeads = 4; +constexpr int32_t kGlmSparseMlaHeadDimQk = 576; +constexpr int32_t kGlmSparseMlaHeadDimV = 512; +constexpr int32_t kGlmSparseMlaMaxTopK = 2048; +constexpr int32_t kGlmSparseMlaVHeadDim = 256; + +struct GlmSparseMlaFlashDecodeParams +{ + // [numTokensQ, 4, 576] E4M3. + void const* q; + // KV pool base: pool token row r holds 576 E4M3 values at byte offset r * 576. + void const* kvPool; + // [numTokensQ, indicesStride] selected tokens of each query token (pool rows, or positions with the remap); + // negative entries are masked. + int32_t const* indices; + // [numTokensQ / seqLenQ] KV length of each sequence, this step's tokens included. + int32_t const* seqLensKv; + // Device scalar that maps Q.K^T to the base-2 softmax domain. + float const* scaleSoftmaxLog2; + // Device scalar that scales the attention output. + float const* outputScale; + // [4, 256, 512] BF16 v_b_proj. + void const* vbProj; + // out[t][h * 256 + v] = sum_d O[t][h][d] * vbProj[h][v][d] (BF16, outStride elements between query tokens), where + // O is the attention output rounded to BF16. + void* out; + int64_t outStride; + int32_t numTokensQ; + // Query tokens per sequence: query token t of a sequence attends to the first + // min(seqLensKv - (seqLenQ - 1 - t), topK) of its indices. + int32_t seqLenQ; + int32_t indicesStride; + int32_t topK; + // Optional in-kernel remap of the indices, with the arithmetic of convertReqIndexToGlobal: with remapBlockTable + // set, indices hold token positions in the query token's sequence, and position p selects pool token row + // remapBlockTable[s * remapBlockTableStride + p / remapTokensPerBlock] * remapStrideFactor + // + p % remapTokensPerBlock + remapLayerOffset, where s = tokenQ / seqLenQ is the sequence (the row of seqLensKv + // and of the block table). A negative position, a block at or beyond remapMaxBlocks or a negative table entry + // masks the entry. + int32_t const* remapBlockTable; + int64_t remapBlockTableStride; + int32_t remapMaxBlocks; + int32_t remapTokensPerBlock; + int32_t remapStrideFactor; + int32_t remapLayerOffset; + // With the remap, true when the indices were written by an earlier layer's top-k (a layer that shares it): the + // kernel then reads the indices, the block table and the KV length before its grid-dependency wait and gathers the + // rows of the positions before this step's tokens (position < seqLensKv - seqLenQ, appended by earlier steps) + // there too; the rows of this step's tokens, written by the KV append it depends on, after the wait. + bool gatherBeforeWait; +}; + +// Generation-only static-token sparse MLA with E4M3 Q and KV and the v_b BMM folded in, as flash decoding over the +// gathered keys. Every query token gets one cluster of 16 CTAs, and CTA s owns indices [128 s, 128 s + 128): +// - each warp gathers 16 token rows with one bulk copy per row onto its own mbarrier; +// - S = Q K^T and P V run as FP16 mma.sync on E4M3 values unpacked exactly to FP16, with fp32 accumulation; the MMAs +// run 8 head columns, the 4 query heads in columns 0-3 and zero Q rows in 4-7, whose outputs are not stored; +// P = exp2(S * scale - max) is stored as E4M3 scaled by 448, as the TRTLLM-Gen generation kernels feed their second +// GEMM, and the query fragments are unpacked once per warp; +// - every CTA sends output slice r of its unnormalized partial output (one 16-column block at a time, during P.V), its +// max and its sum straight into the shared memory of CTA r of the cluster (st.async), and CTA r merges the 16 +// partials of its 32 output columns; +// - CTA r also owns v in [16 r, 16 r + 16) of every head: it all-gathers the row's merged BF16 output (st.async), +// multiplies it with its v_b_proj slice (64 KiB, bulk-copied once every token row has landed) on mma.sync BF16, and +// writes those 64 values of out. +// Requires the SM100 family, numTokensQ * 16 CTAs of 170 KiB shared memory and clusters of 16 CTAs. PDL: the kernel +// reads global memory only after its grid-dependency wait (except the gatherBeforeWait reads above) and lets its +// dependents launch once its partial outputs are sent. +void invokeGlmSparseMlaFlashDecode(GlmSparseMlaFlashDecodeParams const& params, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 4f925e421592..16b2633a737f 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -75,7 +75,8 @@ add_th_common_family( convertReqIndexToGlobalOp.cpp nvfp4MlaKvCacheGatherOp.cpp trtllmGenQKVProcessOp.cpp - glmMlaKvHalfAppendOp.cpp) + glmMlaKvHalfAppendOp.cpp + glmSparseMlaFlashDecodeOp.cpp) add_th_common_family( th_common_gemm_quant diff --git a/cpp/tensorrt_llm/thop/glmSparseMlaFlashDecodeOp.cpp b/cpp/tensorrt_llm/thop/glmSparseMlaFlashDecodeOp.cpp new file mode 100644 index 000000000000..a625400fdc56 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmSparseMlaFlashDecodeOp.cpp @@ -0,0 +1,177 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.h" +#include "tensorrt_llm/runtime/torchUtils.h" + +#include + +#include +#include +#include + +namespace th = torch; +namespace tk = tensorrt_llm::kernels::glm_moe_dsa; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kInt32Max = std::numeric_limits::max(); +// The grid spans the query tokens in its y dimension. +constexpr int64_t kMaxQueryTokens = 65535; +constexpr int64_t kBulkCopyAlignment = 16; + +bool isE4m3Storage(th::Tensor const& t) +{ + return t.scalar_type() == th::kFloat8_e4m3fn || t.scalar_type() == th::kByte; +} + +bool isBulkCopyAligned(th::Tensor const& t) +{ + return reinterpret_cast(t.data_ptr()) % kBulkCopyAlignment == 0; +} + +} // namespace + +// Generation sparse MLA of GLM-5.2 at 4 query heads per rank, with the v_b BMM folded in: writes +// output[t][h * 256 + v] = sum_d bf16(attention(q[t, h]))[d] * v_b_proj[h][v][d] for the first q.size(0) rows of +// output. With block_table, indices hold token positions that the kernel maps to pool rows as +// convert_req_index_to_global does (block_table[s][p / tokens_per_block] * page_index_scale * tokens_per_block +// + p % tokens_per_block + layer_offset * tokens_per_block); without it, indices hold pool rows. +void glmSparseMlaFlashDecode(th::Tensor const& q, th::Tensor const& kvPool, th::Tensor const& indices, + th::Tensor const& seqLensKv, th::Tensor const& mlaBmm1Scale, th::Tensor const& mlaBmm2Scale, + th::Tensor const& vBProj, th::Tensor& output, int64_t seqLenQ, int64_t topK, + std::optional const& blockTable, int64_t tokensPerBlock, int64_t pageIndexScale, int64_t layerOffset, + bool gatherBeforeWait) +{ + TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "glm_sparse_mla_flash_decode needs an SM100-family GPU"); + TORCH_CHECK(q.is_cuda() && kvPool.is_cuda() && indices.is_cuda() && seqLensKv.is_cuda() && mlaBmm1Scale.is_cuda() + && mlaBmm2Scale.is_cuda() && vBProj.is_cuda() && output.is_cuda(), + "glm_sparse_mla_flash_decode takes CUDA tensors"); + auto const device = q.device(); + TORCH_CHECK(kvPool.device() == device && indices.device() == device && seqLensKv.device() == device + && mlaBmm1Scale.device() == device && mlaBmm2Scale.device() == device && vBProj.device() == device + && output.device() == device, + "glm_sparse_mla_flash_decode takes tensors on one device"); + + TORCH_CHECK(isE4m3Storage(q), "q must hold E4M3 values (float8_e4m3fn or uint8)"); + TORCH_CHECK(q.dim() == 3 && q.size(1) == tk::kGlmSparseMlaNumHeads && q.size(2) == tk::kGlmSparseMlaHeadDimQk + && q.is_contiguous() && isBulkCopyAligned(q), + "q must be a contiguous, 16-byte aligned [tokens, ", tk::kGlmSparseMlaNumHeads, ", ", + tk::kGlmSparseMlaHeadDimQk, "] tensor"); + int64_t const numTokens = q.size(0); + TORCH_CHECK(numTokens <= kMaxQueryTokens, "glm_sparse_mla_flash_decode takes at most ", kMaxQueryTokens, + " query tokens, got ", numTokens); + TORCH_CHECK(seqLenQ > 0 && seqLenQ <= kInt32Max && numTokens % seqLenQ == 0, + "seq_len_q must be positive and divide the query tokens; got ", seqLenQ, " for ", numTokens, " tokens"); + int64_t const numSeqs = numTokens / seqLenQ; + + TORCH_CHECK(isE4m3Storage(kvPool), "kv_pool must hold E4M3 values (float8_e4m3fn or uint8)"); + TORCH_CHECK(kvPool.is_contiguous() && isBulkCopyAligned(kvPool) && kvPool.numel() % tk::kGlmSparseMlaHeadDimQk == 0, + "kv_pool must be a contiguous, 16-byte aligned pool of ", tk::kGlmSparseMlaHeadDimQk, "-byte token rows"); + TORCH_CHECK(kvPool.numel() / tk::kGlmSparseMlaHeadDimQk <= kInt32Max, "kv_pool rows must be int32-addressable"); + + TORCH_CHECK(topK > 0 && topK <= tk::kGlmSparseMlaMaxTopK, "top_k must be in [1, ", tk::kGlmSparseMlaMaxTopK, + "], got ", topK); + TORCH_CHECK(indices.scalar_type() == th::kInt32 && indices.dim() == 2 && indices.size(0) == numTokens + && indices.size(1) >= topK && indices.stride(1) == 1 && indices.stride(0) >= topK + && indices.stride(0) <= kInt32Max, + "indices must be an int32 [tokens, >= top_k] tensor with unit column stride"); + TORCH_CHECK(seqLensKv.scalar_type() == th::kInt32 && seqLensKv.dim() == 1 && seqLensKv.size(0) >= numSeqs + && seqLensKv.is_contiguous(), + "seq_lens_kv must be a contiguous int32 tensor with one entry per sequence"); + TORCH_CHECK(mlaBmm1Scale.scalar_type() == th::kFloat32 && mlaBmm1Scale.is_contiguous() && mlaBmm1Scale.numel() >= 2, + "mla_bmm1_scale must hold two contiguous float32 values"); + TORCH_CHECK(mlaBmm2Scale.scalar_type() == th::kFloat32 && mlaBmm2Scale.numel() >= 1, + "mla_bmm2_scale must hold a float32 value"); + + TORCH_CHECK(vBProj.scalar_type() == th::kBFloat16 && vBProj.dim() == 3 + && vBProj.size(0) == tk::kGlmSparseMlaNumHeads && vBProj.size(1) == tk::kGlmSparseMlaVHeadDim + && vBProj.size(2) == tk::kGlmSparseMlaHeadDimV && vBProj.is_contiguous() && isBulkCopyAligned(vBProj), + "v_b_proj must be a contiguous, 16-byte aligned BF16 [", tk::kGlmSparseMlaNumHeads, ", ", + tk::kGlmSparseMlaVHeadDim, ", ", tk::kGlmSparseMlaHeadDimV, "] tensor"); + TORCH_CHECK(output.scalar_type() == th::kBFloat16 && output.dim() == 2 && output.size(0) >= numTokens + && output.size(1) == tk::kGlmSparseMlaNumHeads * tk::kGlmSparseMlaVHeadDim && output.stride(1) == 1, + "output must be a BF16 [>= tokens, ", tk::kGlmSparseMlaNumHeads * tk::kGlmSparseMlaVHeadDim, + "] tensor with unit column stride"); + + tk::GlmSparseMlaFlashDecodeParams params{}; + params.q = q.data_ptr(); + params.kvPool = kvPool.data_ptr(); + params.indices = indices.data_ptr(); + params.seqLensKv = seqLensKv.data_ptr(); + params.scaleSoftmaxLog2 = mlaBmm1Scale.data_ptr() + 1; + params.outputScale = mlaBmm2Scale.data_ptr(); + params.vbProj = vBProj.data_ptr(); + params.out = output.data_ptr(); + params.outStride = output.stride(0); + params.numTokensQ = static_cast(numTokens); + params.seqLenQ = static_cast(seqLenQ); + params.indicesStride = static_cast(indices.stride(0)); + params.topK = static_cast(topK); + params.gatherBeforeWait = gatherBeforeWait; + if (blockTable.has_value()) + { + auto const& table = blockTable.value(); + TORCH_CHECK(table.is_cuda() && table.device() == device && table.scalar_type() == th::kInt32 + && table.dim() == 2 && table.stride(1) == 1 && table.size(0) >= numSeqs && table.size(1) > 0 + && table.size(1) <= kInt32Max && table.stride(0) >= table.size(1), + "block_table must be an int32 CUDA [sequences, max_blocks] tensor with unit column stride"); + TORCH_CHECK(tokensPerBlock > 0 && pageIndexScale > 0 && layerOffset >= 0, "invalid block_table geometry: ", + "tokens_per_block ", tokensPerBlock, ", page_index_scale ", pageIndexScale, ", layer_offset ", layerOffset); + int64_t const strideFactor = pageIndexScale * tokensPerBlock; + int64_t const rowOffset = layerOffset * tokensPerBlock; + TORCH_CHECK(strideFactor <= kInt32Max && rowOffset <= kInt32Max, "block_table rows must be int32-addressable"); + params.remapBlockTable = table.data_ptr(); + params.remapBlockTableStride = table.stride(0); + params.remapMaxBlocks = static_cast(table.size(1)); + params.remapTokensPerBlock = static_cast(tokensPerBlock); + params.remapStrideFactor = static_cast(strideFactor); + params.remapLayerOffset = static_cast(rowOffset); + } + else + { + TORCH_CHECK(!gatherBeforeWait, "gather_before_wait needs block_table"); + } + + auto const stream = at::cuda::getCurrentCUDAStream(q.get_device()); + tk::invokeGlmSparseMlaFlashDecode(params, stream); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_sparse_mla_flash_decode(Tensor q, Tensor kv_pool, Tensor indices, Tensor seq_lens_kv, " + "Tensor mla_bmm1_scale, Tensor mla_bmm2_scale, Tensor v_b_proj, Tensor(a!) output, int seq_len_q, int top_k, " + "Tensor? block_table=None, int tokens_per_block=0, int page_index_scale=0, int layer_offset=0, " + "bool gather_before_wait=False) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_sparse_mla_flash_decode", &tensorrt_llm::torch_ext::glmSparseMlaFlashDecode); +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/sparse_mla_flash_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/sparse_mla_flash_decode.py new file mode 100644 index 000000000000..fea70f687292 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/sparse_mla_flash_decode.py @@ -0,0 +1,363 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DSA sparse MLA generation attention at 4 query heads per rank, with the v_b BMM folded in.""" + +from __future__ import annotations + +import math +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +#: Query heads per rank: GLM-5.2's 64 MLA heads over 16 ranks. +NUM_HEADS = 4 +#: The latent row: kv_lora_rank 512 followed by qk_rope_head_dim 64. +HEAD_DIM_QK = 576 +KV_LORA_RANK = 512 +V_HEAD_DIM = 256 +#: The most selected tokens per query token the kernel reads. +MAX_TOP_K = 2048 +#: The kernel splits the top-k slots over 16 CTAs of 128 slots; each split exponentiates +#: against its own max. +_SPLITS = 16 +_SLOTS_PER_SPLIT = MAX_TOP_K // _SPLITS +#: P is quantized to E4M3 after this scale, as the TRTLLM-Gen generation kernels feed their +#: second GEMM. +_P_SCALE = 448.0 + +#: GLM-5.2's KV pool at TP16: 79 attention layers per rank (78 decoder layers and the MTP +#: layer, at local layer 78), one FP8 page of 64 tokens per layer and pool slot. +_GLM_POOL = dict(tokens_per_block=64, num_layers=79, top_k=MAX_TOP_K) +#: 1 / sqrt(qk_nope_head_dim + qk_rope_head_dim) = 1 / sqrt(192 + 64). +GLM_SOFTMAX_SCALE = 1.0 / math.sqrt(256.0) + + +def _cell(**spec) -> dict: + """A cell spec over GLM-5.2's pool geometry. + + ``kv_lens`` holds each sequence's KV length with this step's ``seq_len_q`` tokens; ``remap`` + passes positions and the block table (else global pool rows); ``layer_offset`` is the pool + layer; ``holes`` leaves masked entries inside the rows' valid range. + """ + return {**_GLM_POOL, "gather_before_wait": False, "holes": False, **spec} + + +class GlmSparseMlaFlashDecode(OpWrapper): + """Sparse MLA flash decode over an FP8 paged latent pool, v_b folded into the output. + + Wraps ``torch.ops.trtllm.glm_sparse_mla_flash_decode``. Query token ``t`` (row ``t`` of + ``q``) belongs to sequence ``s = t // seq_len_q`` and attends to the first + ``n_t = min(seq_lens_kv[s] - (seq_len_q - 1 - t % seq_len_q), top_k)`` slots of + ``indices[t]``; a negative entry is masked. With ``block_table`` a slot holds a position + ``p`` of the sequence, read at pool row:: + + block_table[s, p // tokens_per_block] * page_index_scale * tokens_per_block + + p % tokens_per_block + layer_offset * tokens_per_block + + (``convert_req_index_to_global``'s rows; a block past the table's width or a negative + entry masks the slot). Without it, a slot holds the pool row. Every head ``h`` then gets + the softmax over the selected rows of ``(q[t, h] . row) * mla_bmm1_scale[1]`` (base 2), + applied to the rows' first 512 values and scaled by ``mla_bmm2_scale[0]``, rounded to BF16, + and multiplied by ``v_b_proj[h]``: ``output[t, h * 256 + v]``. A row with nothing selected + returns zeros. + + The probabilities enter the second product as E4M3: each split of 128 slots exponentiates + against its own max, and ``exp2(s - max) * 448`` is rounded to E4M3, while the + normalizing sums add the unrounded values. ``reference`` states that rounding. + + Fusion boundary: the index remap, the gather of the selected rows, the softmax, P.V, the + 16-way merge and the v_b product are inside. Selecting the tokens, the KV append, the + query's projection, RoPE and quantization, and the scales are not. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="verify, full-indexer layer: 8 rows, the layer's own top-k remapped in-kernel", + spec=_cell(seq_len_q=8, kv_lens=(20_011,), remap=True, layer_offset=6), + ), + Cell( + why="verify, shared-indexer layer: the leader's top-k, old rows gathered before wait", + spec=_cell( + seq_len_q=8, kv_lens=(20_011,), remap=True, layer_offset=7, gather_before_wait=True + ), + ), + Cell( + why="MTP draft step 0: 8 rows at the MTP layer, the last layer of the pool", + spec=_cell(seq_len_q=8, kv_lens=(20_018,), remap=True, layer_offset=78), + ), + Cell( + why="MTP draft steps 1-6: one row on global pool rows (reused top-k remapped aside)", + spec=_cell(seq_len_q=1, kv_lens=(20_019,), remap=False, layer_offset=78), + ), + Cell( + why="short context: fewer positions than the top-k, -1 padding and masked holes", + spec=_cell( + seq_len_q=8, + kv_lens=(1_500,), + remap=True, + layer_offset=4, + gather_before_wait=True, + holes=True, + ), + ), + Cell( + why="two sequences: per-sequence block-table rows and KV lengths (max_batch_size 2)", + spec=_cell( + seq_len_q=8, + kv_lens=(9_000, 17_003), + remap=True, + layer_offset=11, + gather_before_wait=True, + ), + ), + ) + + note = """ + Metadata consumed: none -- every input is an argument. What the target passes: + - `q`: the [T, 4, 576] E4M3 query that q_b's fused epilogue writes (absorbed q_nope, + roped q_pe); `seq_len_q` the query tokens per sequence (8 on verify and MTP draft + step 0, 1 on draft steps 1-6). + - `kv_pool`: the DSA cache manager's `get_unique_primary_pool()`, the base that + `host_kv_cache_pool_pointers[pool, 0]` names and `convert_req_index_to_global`'s + rows count from. + - With `block_table` (layers that run the indexer, and the shared layers after one): + `indices` the [T, 2048] local top-k positions (the full layer's own; on a shared + layer its full layer's), `block_table` the DSA metadata's `block_table[num_contexts: + num_seqs]` (pool slots, -1 past each sequence), `(page_index_scale, layer_offset)` + `get_primary_pool_page_index_params(local_layer)`, `tokens_per_block` the manager's. + Without it (draft steps 1-6): `indices` are global pool rows. + - `seq_lens_kv`: the generation sequences' KV lengths with this step's tokens + (`kv_lens_cuda_runtime[num_contexts:num_seqs]`). + - `mla_bmm1_scale` / `mla_bmm2_scale`: the buffers `glm_mla_kv_half_append` fills. + + Masked slots load pool row 0 and multiply it by P = 0: a NaN in that row reaches the + output. Outputs are bitwise deterministic. `output` rows past `q.size(0)` are untouched. + + PDL: the kernel waits on its grid dependency before any global read, except with + `gather_before_wait`: then the indices, the block table, the KV lengths and the pool rows + of positions below `seq_lens_kv - seq_len_q` are read before the wait, so their writers + must have completed before this grid launches (step-start host copies, earlier steps, + and an earlier layer's top-k that some kernel ahead in the stream waited on). A layer + whose own top-k writes `indices` must not set it. Dependents may launch once the partial + outputs are sent; they must wait on their grid dependency before reading `output`. + + The test also replays the shared-layer cell from a captured CUDA graph (the op allocates + nothing and syncs nothing) and runs both ops on a DSACacheManagerV2 pool against + `convert_req_index_to_global`. + + Not certified: `top_k` below 2048, more than two sequences, a ragged batch (present it + token-major: `seq_len_q` 1 with per-row lengths and tables). + """ + + def __call__( + self, + q: torch.Tensor, + kv_pool: torch.Tensor, + indices: torch.Tensor, + seq_lens_kv: torch.Tensor, + mla_bmm1_scale: torch.Tensor, + mla_bmm2_scale: torch.Tensor, + v_b_proj: torch.Tensor, + output: torch.Tensor, + seq_len_q: int, + top_k: int, + block_table: Optional[torch.Tensor] = None, + tokens_per_block: int = 0, + page_index_scale: int = 0, + layer_offset: int = 0, + gather_before_wait: bool = False, + ) -> torch.Tensor: + torch.ops.trtllm.glm_sparse_mla_flash_decode( + q, + kv_pool, + indices, + seq_lens_kv, + mla_bmm1_scale, + mla_bmm2_scale, + v_b_proj, + output, + seq_len_q, + top_k, + block_table, + tokens_per_block, + page_index_scale, + layer_offset, + gather_before_wait, + ) + return output + + def reference( + self, + q: torch.Tensor, + kv_pool: torch.Tensor, + indices: torch.Tensor, + seq_lens_kv: torch.Tensor, + mla_bmm1_scale: torch.Tensor, + mla_bmm2_scale: torch.Tensor, + v_b_proj: torch.Tensor, + output: torch.Tensor, + seq_len_q: int, + top_k: int, + block_table: Optional[torch.Tensor] = None, + tokens_per_block: int = 0, + page_index_scale: int = 0, + layer_offset: int = 0, + gather_before_wait: bool = False, + ) -> torch.Tensor: + """The [T, 1024] BF16 rows the op writes into ``output[:T]``. + + fp32 throughout, with the kernel's three roundings stated: P to E4M3 at scale 448 per + 128-slot split (against the split's max, the sums over the unrounded P), the merged + attention output to BF16 before the v_b product, and the product to BF16. ``output`` + and ``gather_before_wait`` are ignored: one is where the op writes, the other when it + reads. + """ + num_tokens = q.shape[0] + device = q.device + pool_rows = kv_pool.view(torch.uint8).reshape(-1, HEAD_DIM_QK) + + slots = torch.full((num_tokens, MAX_TOP_K), -1, dtype=torch.int64, device=device) + slots[:, :top_k] = indices[:, :top_k].long() + token = torch.arange(num_tokens, device=device) + seq = token // seq_len_q + num_valid = torch.clamp( + seq_lens_kv[seq].long() - (seq_len_q - 1 - token % seq_len_q), max=top_k + ) + if block_table is not None: + position = slots.clamp(min=0) + block = position // tokens_per_block + width = block_table.shape[1] + page = block_table[seq[:, None], block.clamp(max=width - 1)].long() + row = ( + page * (page_index_scale * tokens_per_block) + + position % tokens_per_block + + layer_offset * tokens_per_block + ) + resident = (slots >= 0) & (block < width) & (page >= 0) + else: + row = slots + resident = slots >= 0 + valid = resident & (torch.arange(MAX_TOP_K, device=device)[None, :] < num_valid[:, None]) + row = torch.where(valid, row, torch.zeros_like(row)) + + keys = pool_rows[row].view(torch.float8_e4m3fn).float() + queries = q.view(torch.float8_e4m3fn).float() + with true_fp32_matmul(): + scores = torch.einsum("thd,tjd->thj", queries, keys) * mla_bmm1_scale[1].float() + scores = scores.masked_fill(~valid[:, None, :], float("-inf")) + scores = scores.view(num_tokens, NUM_HEADS, _SPLITS, _SLOTS_PER_SPLIT) + split_max = scores.amax(dim=-1) + empty = split_max == float("-inf") + p = torch.exp2(scores - torch.where(empty, 0.0, split_max)[..., None]) + split_sum = p.sum(dim=-1) + p_e4m3 = (p * _P_SCALE).to(torch.float8_e4m3fn).float() + values = keys[..., :KV_LORA_RANK].reshape( + num_tokens, _SPLITS, _SLOTS_PER_SPLIT, KV_LORA_RANK + ) + with true_fp32_matmul(): + split_out = torch.einsum("thrj,trjd->thrd", p_e4m3, values) + merged_max = split_max.amax(dim=-1, keepdim=True) + alpha = torch.where(empty, 0.0, torch.exp2(split_max - merged_max)) + merged_sum = (alpha * split_sum).sum(dim=-1)[..., None] + merged = (alpha[..., None] * split_out).sum(dim=-2) + attn = torch.where( + merged_sum > 0, + merged / (merged_sum * _P_SCALE) * mla_bmm2_scale.reshape(-1)[0].float(), + 0.0, + ) + attn = attn.to(torch.bfloat16).float() + with true_fp32_matmul(): + out = torch.einsum("thd,hvd->thv", attn, v_b_proj.float()) + return out.reshape(num_tokens, NUM_HEADS * V_HEAD_DIM).to(torch.bfloat16) + + def is_valid( + self, + q: torch.Tensor, + kv_pool: torch.Tensor, + indices: torch.Tensor, + seq_lens_kv: torch.Tensor, + mla_bmm1_scale: torch.Tensor, + mla_bmm2_scale: torch.Tensor, + v_b_proj: torch.Tensor, + output: torch.Tensor, + seq_len_q: int, + top_k: int, + block_table: Optional[torch.Tensor] = None, + tokens_per_block: int = 0, + page_index_scale: int = 0, + layer_offset: int = 0, + gather_before_wait: bool = False, + ) -> None: + """The index and table contents the op reads without checking. + + The op checks every shape, dtype and stride loudly; what it trusts is the contents of + ``indices``, ``seq_lens_kv`` and ``block_table``. ``q``, the scales, ``v_b_proj``, + ``output`` and ``gather_before_wait`` are not read here. + """ + num_tokens = indices.shape[0] + device = indices.device + num_pool_rows = kv_pool.numel() // HEAD_DIM_QK + slots = indices[:, :top_k].long() + token = torch.arange(num_tokens, device=device) + seq = token // seq_len_q + kv_len = seq_lens_kv[seq].long() + num_valid = torch.clamp(kv_len - (seq_len_q - 1 - token % seq_len_q), max=top_k) + unread = torch.arange(slots.shape[1], device=device)[None, :] >= num_valid[:, None] + + # Only the first n_t slots of row t are read: a selection packed elsewhere (a valid + # entry behind them) silently drops tokens. + assert not ((slots >= 0) & unread).any(), ( + "indices holds a selected token past the row's valid count " + "min(kv_len - (seq_len_q - 1 - t), top_k); the op never reads it" + ) + selected = (slots >= 0) & ~unread + if block_table is None: + # Global rows are read as given: one past the pool reads other memory. + assert not (selected & (slots >= num_pool_rows)).any(), ( + f"indices names a pool row past the pool's {num_pool_rows} rows" + ) + return + + # The positions are the query token's own context: a later position reads a row this + # step has not appended, or another request's. + own_position = kv_len - seq_len_q + token % seq_len_q + assert not (selected & (slots > own_position[:, None])).any(), ( + "indices selects a position past the query token's own: the row is not this " + "token's context" + ) + # The layer's rows of a slot lie inside the slot: an offset at or past the page scale + # reads the next slot's pages. + assert 0 <= layer_offset < page_index_scale, ( + f"layer_offset {layer_offset} must be below page_index_scale {page_index_scale}" + ) + # Every block of a sequence's KV must be mapped: a narrower table or a -1 inside the + # sequence masks live tokens instead of failing. + num_seqs = num_tokens // seq_len_q + lens = seq_lens_kv[:num_seqs].long() + blocks = (lens + tokens_per_block - 1) // tokens_per_block + width = block_table.shape[1] + assert bool((blocks <= width).all()), ( + f"block_table has {width} columns; a sequence needs {int(blocks.max())}" + ) + live = torch.arange(width, device=device)[None, :] < blocks[:, None] + table = block_table[:num_seqs].long() + assert not (live & (table < 0)).any(), "block_table has a -1 inside a sequence's KV" + if bool(live.any()): + last_page = int(table[live].max()) * page_index_scale + layer_offset + end_row = (last_page + 1) * tokens_per_block + assert end_row <= num_pool_rows, ( + f"block_table maps rows up to {end_row}, past the pool's {num_pool_rows} rows" + ) + + +sparse_mla_flash_decode = GlmSparseMlaFlashDecode() diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_sparse_mla_flash_decode.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_sparse_mla_flash_decode.py new file mode 100644 index 000000000000..e2fe77249109 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_sparse_mla_flash_decode.py @@ -0,0 +1,563 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM-5.2 tp16 sparse_mla_flash_decode entry. + +The cells build GLM-5.2's TP16 pool geometry directly: FP8 latent rows in 79 layers' 64-token +pages per pool slot, a block table that maps each sequence's blocks to shuffled slots (-1 past +the sequence), and top-k rows laid out as the DSA indexer writes them -- the valid entries +packed in front, this step's visible tokens among them, -1 behind. Every cell is driven through +the entry's guard and gated against the entry's reference. + +Two further tests: a captured CUDA graph replays the shared-layer cell bit for bit, and on a +real DSACacheManagerV2 the KV append writes this step's rows where main's +`convert_req_index_to_global` says they live and the flash decode's in-kernel remap reads +exactly those rows. +""" + +import math +import zlib +from typing import NamedTuple + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.attention.mla_kv_half_append import ( # noqa: E501 + LATENT_OFFSET, + mla_kv_half_append, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.attention.sparse_mla_flash_decode import ( # noqa: E501 + GLM_SOFTMAX_SCALE, + HEAD_DIM_QK, + KV_LORA_RANK, + MAX_TOP_K, + NUM_HEADS, + V_HEAD_DIM, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.attention.sparse_mla_flash_decode import ( # noqa: E501 + sparse_mla_flash_decode as op, +) +from tensorrt_llm._torch.attention.backends.interface import RopeParams +from tensorrt_llm._torch.attention.backends.sparse.dsa.cache_manager import DSACacheManagerV2 +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig, KvCacheConfig +from tensorrt_llm.mapping import Mapping + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or get_sm_version() != 107, + reason="the GLM-5.2 tp16 target's ops are certified on sm_107", +) + +#: log2(e) as an fp32 constant, as the KV append forms mla_bmm1_scale[1]. +_LOG2E = 1.4426950408889634 +#: Pool slots no sequence maps, and block-table columns past every sequence. +_SPARE_SLOTS = 5 +_SPARE_COLUMNS = 3 +#: Output rows past the query tokens, which the op must leave alone. +_SPARE_ROWS = 2 +#: Magnitudes of the random operands: KV rows ~ N(0, 1) as kv_a_layernorm leaves them, the +#: absorbed query ~ N(0, 0.25), v_b ~ N(0, 1 / 512). +_Q_STD = 0.5 +_VB_STD = 1.0 / math.sqrt(KV_LORA_RANK) +#: The relative L2 distance allowed between the reference and an attention whose +#: probabilities are not rounded to E4M3. E4M3 keeps 3 mantissa bits: a rounded probability is +#: off by up to 2**-4 relative (RMS 1.8-3.6 % across a binade), and on outputs that average +#: zero-mean rows the error stays at that relative size: the kernel measured 2.3 % against an +#: fp32 reference, and the bound is 4x that. A reference that mis-masks a slot, +#: drops a split or misses a scale is off by tens of percent. +_E4M3_P_REL_L2 = 0.1 + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """4 ulp of each token row's largest output, 1 ulp relative RMS, in BF16. + + The reference states the kernel's three roundings (P to E4M3 per 128-slot split, the merged + attention output to BF16, the v_b product to BF16), so what is left is fp32 accumulation + order: the QK and P.V sums inside mma.sync, the 16-way merge, exp2f, and the v_b sum. Those + move a value by about 1e-6 relative; that flips an E4M3 probability with probability about + 1e-5 per element (one flip in a cell of 64k probabilities at most), and one flipped + probability moves the outputs of its head by under 1 % of their magnitude (2 ulp of the + element, under 1 ulp of the row's largest); the final BF16 roundings flip by one ulp at + the same rate. A mirror with a wrong merge weight, a wrong split boundary or a dropped + head is off by many percent. + """ + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=1.0) + + +def _fill_e4m3_(dst: torch.Tensor, gen: torch.Generator, std: float = 1.0) -> None: + """Fill a contiguous uint8 view with E4M3 values of N(0, std^2) (finite: far below 448).""" + flat = dst.reshape(-1) + chunk = 1 << 26 + for start in range(0, flat.numel(), chunk): + stop = min(start + chunk, flat.numel()) + values = torch.randn(stop - start, generator=gen, device=flat.device) * std + flat[start:stop].copy_(values.to(torch.float8_e4m3fn).view(torch.uint8)) + + +def _select(kv_lens, seq_len_q: int, top_k: int, holes: bool, gen: torch.Generator) -> torch.Tensor: + """[T, top_k] int32 top-k positions as the DSA indexer lays them out. + + Row t of sequence s selects min(own + 1, top_k) positions out of 0 .. own (own = its + position): every token of this step it can see, the rest drawn from the earlier steps', + shuffled, then -1. With `holes`, about 1/64 of the selected slots become -1 in place. + """ + rows = [] + for kv_len in kv_lens: + first_new = kv_len - seq_len_q + for t in range(seq_len_q): + own = first_new + t + count = min(own + 1, top_k) + fresh = torch.arange(first_new, own + 1, dtype=torch.int64) + num_old = count - fresh.numel() + old = torch.randperm(first_new, generator=gen)[:num_old] + chosen = torch.cat((fresh, old)) + chosen = chosen[torch.randperm(chosen.numel(), generator=gen)] + if holes: + drop = torch.randperm(count, generator=gen)[: max(1, count // 64)] + chosen[drop] = -1 + row = torch.full((top_k,), -1, dtype=torch.int32) + row[:count] = chosen.to(torch.int32) + rows.append(row) + return torch.stack(rows) + + +def _global_rows( + positions: torch.Tensor, + table: torch.Tensor, + seq_len_q: int, + tokens_per_block: int, + page_index_scale: int, + layer_offset: int, +) -> torch.Tensor: + """The pool rows of `positions` with convert_req_index_to_global's arithmetic (-1 kept).""" + num_tokens = positions.shape[0] + seq = torch.arange(num_tokens, device=positions.device) // seq_len_q + position = positions.long() + block = position.clamp(min=0) // tokens_per_block + page = table[seq[:, None], block.clamp(max=table.shape[1] - 1)].long() + rows = ( + page * page_index_scale * tokens_per_block + + layer_offset * tokens_per_block + + position.clamp(min=0) % tokens_per_block + ) + return torch.where((position >= 0) & (page >= 0), rows, -1).to(torch.int32) + + +class _Case(NamedTuple): + kwargs: dict + num_tokens: int + global_rows: torch.Tensor + + +def _build(spec: dict, seed: int) -> _Case: + """Turn a cell spec into the op's arguments, over a pool of the spec's geometry.""" + gen_cpu = torch.Generator().manual_seed(seed) + gen = torch.Generator(device="cuda").manual_seed(seed) + tokens_per_block = spec["tokens_per_block"] + num_layers = spec["num_layers"] + layer_offset = spec["layer_offset"] + seq_len_q = spec["seq_len_q"] + top_k = spec["top_k"] + kv_lens = list(spec["kv_lens"]) + num_tokens = len(kv_lens) * seq_len_q + + blocks = [-(-kv_len // tokens_per_block) for kv_len in kv_lens] + num_slots = sum(blocks) + _SPARE_SLOTS + order = torch.randperm(num_slots, generator=gen_cpu).to(torch.int32) + table = torch.full((len(kv_lens), max(blocks) + _SPARE_COLUMNS), -1, dtype=torch.int32) + start = 0 + for s, num_blocks in enumerate(blocks): + table[s, :num_blocks] = order[start : start + num_blocks] + start += num_blocks + + pool = torch.empty( + num_slots, num_layers, tokens_per_block * HEAD_DIM_QK, dtype=torch.uint8, device="cuda" + ) + _fill_e4m3_(pool, gen) + positions = _select(kv_lens, seq_len_q, top_k, spec["holes"], gen_cpu) + global_rows = _global_rows( + positions, table, seq_len_q, tokens_per_block, num_layers, layer_offset + ) + + scale = torch.tensor(GLM_SOFTMAX_SCALE, dtype=torch.float32) + kwargs = dict( + q=( + torch.randn(num_tokens, NUM_HEADS, HEAD_DIM_QK, generator=gen, device="cuda") * _Q_STD + ).to(torch.float8_e4m3fn), + kv_pool=pool.view(torch.float8_e4m3fn), + indices=(positions if spec["remap"] else global_rows).cuda(), + seq_lens_kv=torch.tensor(kv_lens, dtype=torch.int32, device="cuda"), + mla_bmm1_scale=torch.stack((scale, scale * torch.tensor(_LOG2E))).cuda(), + mla_bmm2_scale=torch.ones(1, dtype=torch.float32, device="cuda"), + v_b_proj=( + torch.randn(NUM_HEADS, V_HEAD_DIM, KV_LORA_RANK, generator=gen, device="cuda") + * _VB_STD + ).to(torch.bfloat16), + output=torch.full( + (num_tokens + _SPARE_ROWS, NUM_HEADS * V_HEAD_DIM), + float("nan"), + dtype=torch.bfloat16, + device="cuda", + ), + seq_len_q=seq_len_q, + top_k=top_k, + ) + if spec["remap"]: + kwargs.update( + block_table=table.cuda(), + tokens_per_block=tokens_per_block, + page_index_scale=num_layers, + layer_offset=layer_offset, + gather_before_wait=spec["gather_before_wait"], + ) + return _Case(kwargs=kwargs, num_tokens=num_tokens, global_rows=global_rows.cuda()) + + +def _unrounded_attention(case: _Case) -> torch.Tensor: + """[T, 1024] fp32: softmax attention over the selected rows without any intermediate + rounding, times v_b. An independent statement of what the op approximates.""" + kw = case.kwargs + num_tokens, seq_len_q, top_k = case.num_tokens, kw["seq_len_q"], kw["top_k"] + rows = case.global_rows.long() + token = torch.arange(num_tokens, device="cuda") + kv_len = kw["seq_lens_kv"][token // seq_len_q].long() + num_valid = torch.clamp(kv_len - (seq_len_q - 1 - token % seq_len_q), max=top_k) + valid = (rows >= 0) & (torch.arange(top_k, device="cuda")[None, :] < num_valid[:, None]) + pool_rows = kw["kv_pool"].view(torch.uint8).reshape(-1, HEAD_DIM_QK) + keys = pool_rows[rows.clamp(min=0)].view(torch.float8_e4m3fn).float().double() + queries = kw["q"].float().double() + scores = torch.einsum("thd,tjd->thj", queries, keys) * GLM_SOFTMAX_SCALE + scores = scores.masked_fill(~valid[:, None, :], float("-inf")) + weights = torch.nan_to_num(torch.softmax(scores, dim=-1), nan=0.0) + attn = torch.einsum("thj,tjd->thd", weights, keys[..., :KV_LORA_RANK]) + out = torch.einsum("thd,hvd->thv", attn, kw["v_b_proj"].double()) + return out.reshape(num_tokens, NUM_HEADS * V_HEAD_DIM).float() + + +def _seed(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by name + certifying(op) + case = _build(cell.spec, _seed(cell.why)) + output = case.kwargs["output"] + op(**case.kwargs) + torch.cuda.synchronize() + ref = op.reference(**case.kwargs) + + # The entry's reference is attention: within the E4M3-probability band of an unrounded + # one. Checked before the kernel gate, so a mirror that drifted from attention fails as + # such rather than as a kernel mismatch. + unrounded = _unrounded_attention(case) + rel_l2 = ((ref.float() - unrounded).norm() / unrounded.norm()).item() + assert rel_l2 <= _E4M3_P_REL_L2, f"reference is {rel_l2:.3%} from unrounded attention" + + _compare(output[: case.num_tokens], ref) + untouched = output[case.num_tokens :].float() + assert torch.isnan(untouched).all(), "rows past the tokens were written" + + +def test_capture_and_replay_match_the_eager_call() -> None: + """The shared-layer call (early gather) captured in a CUDA graph replays bit for bit, and + two eager calls agree bit for bit.""" + cell = op.CELLS[1] + assert cell.spec["gather_before_wait"] + case = _build(cell.spec, _seed(cell.why)) + output = case.kwargs["output"] + op(**case.kwargs) + torch.cuda.synchronize() + eager = output.clone() + output.fill_(float("nan")) + op(**case.kwargs) + torch.cuda.synchronize() + assert torch.equal(output[: case.num_tokens], eager[: case.num_tokens]), "not deterministic" + + graph = torch.cuda.CUDAGraph() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + with torch.cuda.graph(graph, stream=stream): + op(**case.kwargs) + torch.cuda.current_stream().wait_stream(stream) + for _ in range(3): + output.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(output[: case.num_tokens], eager[: case.num_tokens]) + + +def _guard_case(cell_index: int) -> _Case: + cell = op.CELLS[cell_index] + return _build(cell.spec, _seed(cell.why)) + + +def test_guard_refuses_a_selection_past_the_valid_count() -> None: + """A selected token behind the row's valid count is never read.""" + case = _guard_case(4) # short context: rows hold fewer valid slots than the top-k + seq_len_q = case.kwargs["seq_len_q"] + first_unread = int(case.kwargs["seq_lens_kv"][0]) - (seq_len_q - 1) # row 0's valid count + indices = case.kwargs["indices"] + assert first_unread < indices.shape[1] + indices[0, first_unread] = 0 + with pytest.raises(AssertionError, match="past the row's valid count"): + with validating(op): + op(**case.kwargs) + + +def test_guard_refuses_a_position_past_the_token_itself() -> None: + """A position after the query token's own reads a row it may not see.""" + case = _guard_case(0) + kv_len = int(case.kwargs["seq_lens_kv"][0]) + case.kwargs["indices"][0, 0] = kv_len - 1 # row 0 sits at kv_len - 8 + with pytest.raises(AssertionError, match="past the query token's own"): + with validating(op): + op(**case.kwargs) + + +def test_guard_refuses_a_table_narrower_than_the_kv() -> None: + """A table cut short masks live tokens instead of failing.""" + case = _guard_case(0) + kv_len = int(case.kwargs["seq_lens_kv"][0]) + needed = -(-kv_len // case.kwargs["tokens_per_block"]) + case.kwargs["block_table"] = case.kwargs["block_table"][:, : needed - 1] + with pytest.raises(AssertionError, match="columns; a sequence needs"): + with validating(op): + op(**case.kwargs) + + +def test_guard_refuses_a_hole_in_the_table() -> None: + """A -1 inside a sequence's blocks masks its tokens instead of failing.""" + case = _guard_case(0) + case.kwargs["block_table"][0, 3] = -1 + with pytest.raises(AssertionError, match="-1 inside"): + with validating(op): + op(**case.kwargs) + + +def test_guard_refuses_a_layer_offset_outside_the_slot() -> None: + """An offset at the page scale reads the next slot's first layer.""" + case = _guard_case(0) + case.kwargs["layer_offset"] = case.kwargs["page_index_scale"] + with pytest.raises(AssertionError, match="below page_index_scale"): + with validating(op): + op(**case.kwargs) + + +def test_guard_refuses_a_global_row_past_the_pool() -> None: + """Global rows are read as given.""" + case = _guard_case(3) # draft steps 1-6: global rows + assert "block_table" not in case.kwargs + num_rows = case.kwargs["kv_pool"].numel() // HEAD_DIM_QK + case.kwargs["indices"][0, 0] = num_rows + with pytest.raises(AssertionError, match="past the pool"): + with validating(op): + op(**case.kwargs) + + +class _DsaEnv: + """A DSACacheManagerV2 over GLM-5.2's latent rows (FP8, 576 per token, 64-token pages).""" + + TOKENS_PER_BLOCK = 64 + + def __init__(self, num_layers: int, max_batch_size: int, max_seq_len: int) -> None: + self.max_batch_size = max_batch_size + self.manager = DSACacheManagerV2( + kv_cache_config=KvCacheConfig( + max_tokens=max_batch_size * (max_seq_len + self.TOKENS_PER_BLOCK), + enable_block_reuse=False, + ), + kv_cache_type=CacheType.SELFKONLY, + num_layers=num_layers, + num_kv_heads=1, + head_dim=HEAD_DIM_QK, + tokens_per_block=self.TOKENS_PER_BLOCK, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=DataType.FP8, + sparse_attention_config=DeepSeekSparseAttentionConfig( + index_head_dim=128, index_n_heads=32, index_topk=MAX_TOP_K + ), + ) + rope = RopeParams( + dim=HEAD_DIM_QK - KV_LORA_RANK, + theta=8_000_000.0, + max_positions=max_seq_len, + duplicate_data=True, + ) + _, self.rotary_cos_sin = rope.create_rope_const_params() + + def generation_step(self, cached: list, seq_len_q: int) -> TrtllmAttentionMetadata: + """Requests 0 .. n-1 with `cached` tokens each, grown by one step of `seq_len_q`.""" + request_ids = list(range(len(cached))) + self.manager.add_dummy_requests(request_ids, token_nums=cached) + for rid in request_ids: + cache = self.manager.kv_cache_map[rid] + assert cache.resize(cache.capacity + seq_len_q) + metadata = TrtllmAttentionMetadata( + max_num_requests=self.max_batch_size, + max_num_tokens=8192, + kv_cache_manager=self.manager, + ) + metadata.seq_lens = torch.tensor([seq_len_q] * len(cached), dtype=torch.int) + metadata.num_contexts = 0 + metadata.request_ids = request_ids + metadata.prompt_lens = cached + metadata.kv_cache_params = KVCacheParams( + use_cache=True, num_cached_tokens_per_seq=cached + ) + metadata.prepare() + return metadata + + def block_table(self, kv_lens: list) -> torch.Tensor: + """The DSA metadata's generation block table: pool slots, -1 past each sequence.""" + slots = self.manager.get_pool_block_indices( + len(kv_lens), request_ids=list(range(len(kv_lens))), num_contexts=0 + ) + blocks = torch.tensor([-(-n // self.TOKENS_PER_BLOCK) for n in kv_lens]) + width = int(blocks.max()) + table = slots[:, :width].to(torch.int32).clone() + table[torch.arange(width)[None, :] >= blocks[:, None]] = -1 + return table.cuda() + + def shutdown(self) -> None: + self.manager.shutdown() + + +def _convert( + block_table: torch.Tensor, positions: torch.Tensor, seq_len_q: int, page_params, tpb: int +) -> torch.Tensor: + """Main's pool rows of `positions` ([T, k] int32): convert_req_index_to_global.""" + page_index_scale, layer_offset = page_params + num_tokens = positions.shape[0] + req_idx = (torch.arange(num_tokens, device="cuda") // seq_len_q).to(torch.int32) + return torch.ops.trtllm.convert_req_index_to_global( + req_idx, + block_table, + positions, + tpb, + positions.shape[1], + page_index_scale * tpb, + layer_offset, + ) + + +def test_reads_the_rows_the_kv_append_writes_in_the_dsa_pool() -> None: + """Both ops against main's DSA cache layout. + + On a DSACacheManagerV2 pool, glm_mla_kv_half_append appends this step's 8 rows of one layer + through the attention metadata's block offsets; they must land where main's + convert_req_index_to_global places those positions. The flash decode then remaps the step's + top-k positions in the kernel, with and without the early gather, and must read exactly the + rows convert_req_index_to_global names: its output is gated against the reference run on + those global rows. + """ + seq_len_q, layer = 8, 2 + cached = [5_000] + kv_lens = [c + seq_len_q for c in cached] + num_tokens = len(cached) * seq_len_q + gen_cpu = torch.Generator().manual_seed(1) + gen = torch.Generator(device="cuda").manual_seed(1) + env = _DsaEnv(num_layers=4, max_batch_size=2, max_seq_len=8192) + try: + metadata = env.generation_step(cached, seq_len_q) + tpb = env.TOKENS_PER_BLOCK + pool = env.manager.get_unique_primary_pool() + pool_rows = pool.view(torch.uint8).reshape(-1, HEAD_DIM_QK) + _fill_e4m3_(pool_rows, gen) + block_table = env.block_table(kv_lens) + page_params = env.manager.get_primary_pool_page_index_params(layer) + + fused_a = torch.randn(num_tokens, 2624, generator=gen, device="cuda").to(torch.bfloat16) + latent = fused_a[:, LATENT_OFFSET : LATENT_OFFSET + HEAD_DIM_QK] + kv_norm_weight = ( + 1.0 + 0.1 * torch.randn(KV_LORA_RANK, generator=gen, device="cuda") + ).to(torch.bfloat16) + bmm1 = torch.full((2,), float("nan"), dtype=torch.float32, device="cuda") + bmm2 = torch.full((1,), float("nan"), dtype=torch.float32, device="cuda") + append_args = ( + latent, + kv_norm_weight, + 1e-5, + env.rotary_cos_sin, + metadata.kv_lens_cuda_runtime, + 0, + seq_len_q, + metadata.kv_cache_block_offsets, + env.manager.kv_cache_pool_pointers, + env.manager.kv_cache_pool_mapping, + layer, + tpb, + env.manager.max_seq_len, + bmm1, + bmm2, + GLM_SOFTMAX_SCALE, + True, + ) + with validating(mla_kv_half_append): + mla_kv_half_append(*append_args) + torch.cuda.synchronize() + want_rows, positions, want_bmm1, want_bmm2 = mla_kv_half_append.reference(*append_args) + fresh = _convert( + block_table, positions.to(torch.int32)[:, None], seq_len_q, page_params, tpb + ) + got_rows = pool_rows[fresh[:, 0].long()].view(torch.float8_e4m3fn) + torch.testing.assert_close(got_rows.float(), want_rows.float(), rtol=2**-3, atol=2**-9) + assert torch.equal(bmm1, want_bmm1) and torch.equal(bmm2, want_bmm2) + + indices = _select(kv_lens, seq_len_q, MAX_TOP_K, False, gen_cpu).cuda() + global_rows = _convert(block_table, indices, seq_len_q, page_params, tpb) + q = ( + torch.randn(num_tokens, NUM_HEADS, HEAD_DIM_QK, generator=gen, device="cuda") * _Q_STD + ).to(torch.float8_e4m3fn) + v_b_proj = ( + torch.randn(NUM_HEADS, V_HEAD_DIM, KV_LORA_RANK, generator=gen, device="cuda") + * _VB_STD + ).to(torch.bfloat16) + output = torch.empty( + num_tokens, NUM_HEADS * V_HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + ref = op.reference( + q, + pool, + global_rows, + metadata.kv_lens_cuda_runtime, + bmm1, + bmm2, + v_b_proj, + output, + seq_len_q, + MAX_TOP_K, + ) + for gather_before_wait in (False, True): + output.fill_(float("nan")) + with validating(op): + op( + q, + pool, + indices, + metadata.kv_lens_cuda_runtime, + bmm1, + bmm2, + v_b_proj, + output, + seq_len_q, + MAX_TOP_K, + block_table, + tpb, + page_params[0], + page_params[1], + gather_before_wait, + ) + torch.cuda.synchronize() + _compare(output, ref) + finally: + env.shutdown() From 0910c8d4600326eb2d2366bd7164033ffd4e2f39 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:18:52 -0700 Subject: [PATCH 21/53] Run the GLM-5.2 decode step's attention core on the q_b, KV-half and flash-decode ops For up to 8 rows the decode step's attention runs q_b with k_b absorbed and the q_a norm folded (built once from the loaded weights), whose epilogue writes the RoPE'd FP8 query, beside the KV half on the aux stream (kv_a norm, K RoPE, FP8 append and the attention scales in one kernel), and then the sparse MLA flash decode at 4 heads with v_b folded in over the built-in indexer's selection. The built-in indexer projection, K-cache update and top-k selection are unchanged. Signed-off-by: Dongfeng Yu --- .../glm_5_2_nvfp4__sm_107__tp16/modeling.py | 170 ++++++++++++++++-- 1 file changed, 154 insertions(+), 16 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py index 8ad9e6836ad8..013d8b32b6ff 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py @@ -22,10 +22,15 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( fp4_quantize, ) -from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + AttentionInputType, + AttentionMetadata, +) from tensorrt_llm._torch.attention.backends.sparse.dsa.module import ( forward_generation_sparse_attn, ) +from tensorrt_llm._torch.attention.backends.sparse.dsa.params import DSABackendForwardArgs from tensorrt_llm._torch.distributed import AllReduceFusionOp, AllReduceParams from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3ForCausalLM, DeepseekV3Model @@ -42,12 +47,15 @@ from tensorrt_llm._torch.utils import EventType from . import weights as _weights +from .ops.attention.mla_kv_half_append import mla_kv_half_append +from .ops.attention.sparse_mla_flash_decode import sparse_mla_flash_decode from .ops.embedding.fused_embedding import fused_embedding from .ops.gemm.dense_mlp_small_m import INTER_PER_RANK, dense_mlp_small_m from .ops.gemm.fused_a_gemm import MAX_TOKENS as FUSED_A_MAX_TOKENS from .ops.gemm.fused_a_gemm import OUT_FEATURES as FUSED_A_OUT_FEATURES from .ops.gemm.fused_a_gemm import fused_a_gemm from .ops.gemm.low_m_gemm import low_m_gemm +from .ops.gemm.q_b_fp8_rope import q_b_fp8_rope from .ops.moe.routed_experts import ( FUSED_ROUTER_MAX_TOKENS, make_fused_router_workspace, @@ -102,10 +110,14 @@ class DecodeTarget: following norm; the experts finalize before theirs (16 ranks span nodes). - The last all-reduce applies the model's final norm (``next_layer_layernorm``). - Attention is MLA with the DSA indexer on the generation rows: q_a, the - compressed KV and its rope part come from one GEMM over the stored - kv_a_proj_with_mqa weight (up to 8 rows), and the built-in modules compute - the rest. + Attention is MLA with the DSA indexer on the generation rows. For up to 8 + rows: q_a, the compressed KV and its rope part come from one GEMM over the + stored kv_a_proj_with_mqa weight; q_b with k_b absorbed and the q_a norm + folded writes the RoPE'd FP8 query in its epilogue, while on the aux stream + one kernel normalizes, ropes and appends the KV half to the FP8 cache; the + built-in indexer selects the 2048 rows; and the flash decode attends over + them at 4 heads with v_b folded in, ahead of o_proj. A larger batch takes the + built-in modules' path. The MoE block over up to 8 rows runs its shared expert as two split-K GEMMs around SwiGLU on the aux stream, and its routed experts in one @@ -167,12 +179,41 @@ def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: and not backend.bias and backend.w3_w1_weight.shape[-1] * 2 == config.hidden_size ), f"layer {layer.layer_idx}: routed experts are not unpadded TRTLLM-Gen NVFP4" + self.index_topk = config.index_topk + # q_b with k_b absorbed and the q_a norm weight folded, per layer: for head h, rows + # [h * 576, h * 576 + 512) are k_b_trans[h] @ W_q_b[h, :qk_nope] and the next 64 rows + # W_q_b[h, qk_nope:], every column times gamma_q_a; FP64, rounded to BF16 once. + self._q_b_absorbed = {} + for layer in core.layers[: core.num_hidden_layers]: + mla = layer.self_attn + assert ( + getattr(mla.mqa, "has_fp8_kv_cache", False) + and mla.kv_cache_dtype != "fp8_ds_mla" + and mla.k_b_proj_trans.dtype == torch.bfloat16 + and mla.v_b_proj.dtype == torch.bfloat16 + and mla.v_b_proj.is_contiguous() + and mla.q_b_proj.weight.dtype == torch.bfloat16 + ), f"layer {layer.layer_idx}: unexpected MLA weights or KV cache layout" + qk_head_dim = mla.qk_nope_head_dim + mla.qk_rope_head_dim + w = mla.q_b_proj.weight.view(mla.num_heads_tp, qk_head_dim, mla.q_lora_rank).double() + nope = torch.bmm(mla.k_b_proj_trans.double(), w[:, : mla.qk_nope_head_dim]) + absorbed = torch.cat([nope, w[:, mla.qk_nope_head_dim :]], 1) + absorbed = absorbed * mla.q_a_layernorm.weight.double() + self._q_b_absorbed[layer.layer_idx] = ( + absorbed.to(torch.bfloat16).reshape(-1, mla.q_lora_rank).contiguous() + ) + device = core.embed_tokens.weight.device + # The attention scales the KV-half kernel writes and the flash decode reads. + self._bmm1_scale = torch.empty(2, dtype=torch.float32, device=device) + self._bmm2_scale = torch.empty(1, dtype=torch.float32, device=device) + self._unit_scale = torch.ones(1, dtype=torch.float32, device=device) + self._kv_pool = None + moe_layers = [ layer for layer in core.layers[: core.num_hidden_layers] if not isinstance(layer.mlp, GatedMLP) ] - device = core.embed_tokens.weight.device # One exchange workspace per router call site: its launches must run in stream order. self._router_workspace = { layer.layer_idx: make_fused_router_workspace(device) for layer in moe_layers @@ -290,15 +331,113 @@ def shared() -> torch.Tensor: ) return shared_output.add_(routed_output) - @staticmethod - def _attention(layer, position_ids, hidden_states, attn_metadata, all_reduce_params, fused_a): + def _attention(self, layer, position_ids, hidden_states, attn_metadata, all_reduce_params, fused): """MLA with the DSA indexer over generation rows, through o_proj.""" + if not fused: + return self._attention_builtin( + layer, position_ids, hidden_states, attn_metadata, all_reduce_params + ) mla = layer.self_attn + mqa = mla.mqa + manager = attn_metadata.kv_cache_manager num_tokens = attn_metadata.num_tokens - if fused_a: - kv_a = fused_a_gemm(hidden_states, mla.kv_a_proj_with_mqa.weight, []) - else: - kv_a = mla.kv_a_proj_with_mqa(hidden_states) + seq_len_q = num_tokens // attn_metadata.num_seqs + head_dim = mla.kv_lora_rank + mla.qk_rope_head_dim + mqa._ensure_rope_table_size(attn_metadata.max_seq_len) + cos_sin = mqa.rotary_cos_sin + kv_lens = attn_metadata.kv_lens_cuda_runtime + if self._kv_pool is None: + self._kv_pool = manager.get_unique_primary_pool() + + kv_a = fused_a_gemm(hidden_states, mla.kv_a_proj_with_mqa.weight, [])[:num_tokens] + q_a = kv_a[:, : mla.q_lora_rank] + quant_q = hidden_states.new_empty((num_tokens, mla.num_heads_tp * head_dim), dtype=torch.uint8) + + def q_half(): + q_b_fp8_rope( + q_a, + self._q_b_absorbed[layer.layer_idx], + quant_q, + kv_lens, + cos_sin.view(-1, mla.qk_rope_head_dim, 2), + self._unit_scale, + mla.q_a_layernorm.variance_epsilon, + head_dim, + mla.qk_rope_head_dim, + seq_len_q, + "regular", + ) + + def kv_half(): + mla_kv_half_append( + kv_a[:, mla.q_lora_rank :], + mla.kv_a_layernorm.weight, + mla.kv_a_layernorm.variance_epsilon, + cos_sin, + kv_lens, + 0, + seq_len_q, + attn_metadata.kv_cache_block_offsets, + manager.kv_cache_pool_pointers, + manager.kv_cache_pool_mapping, + mqa.get_local_layer_idx(attn_metadata), + manager.tokens_per_block, + attn_metadata.max_seq_len, + self._bmm1_scale, + self._bmm2_scale, + mla.softmax_scale, + ) + + maybe_execute_in_parallel( + q_half, kv_half, mla.ln_events[0], mla.ln_events[1], mla.aux_stream + ) + indexer = mqa.indexer + indexer_intermediates = None + if indexer is not None: + qr = mla.q_a_layernorm(q_a) + indexer_intermediates = [ + t[:num_tokens] + for t in indexer.pre_indexer_proj( + qr, hidden_states[:num_tokens], position_ids[..., :num_tokens] + ) + ] + indexer._update_k_cache(indexer_intermediates[1], indexer_intermediates[2], attn_metadata) + indices, _ = mqa.sparse_attn_predict( + hidden_states[:num_tokens], + None, + attn_metadata, + AttentionForwardArgs( + attention_input_type=AttentionInputType.generation_only, + sparse_backend_args=DSABackendForwardArgs( + indexer_intermediates=indexer_intermediates + ), + ), + ) + output = hidden_states.new_empty( + (hidden_states.shape[0], mla.num_heads_tp * mla.v_head_dim) + ) + sparse_mla_flash_decode( + quant_q.view(num_tokens, mla.num_heads_tp, head_dim), + self._kv_pool, + indices, + kv_lens, + self._bmm1_scale, + self._bmm2_scale, + mla.v_b_proj, + output, + seq_len_q, + self.index_topk, + ) + if indexer is not None: + indexer.maybe_join_prev_topk_copy() + return mla.o_proj(output, all_reduce_params=all_reduce_params, layer_idx=mla.layer_idx) + + @staticmethod + def _attention_builtin(layer, position_ids, hidden_states, attn_metadata, all_reduce_params): + """The built-in modules' MLA with the DSA indexer, for a batch the kernels do not serve.""" + mla = layer.self_attn + num_tokens = attn_metadata.num_tokens + kv_a = mla.kv_a_proj_with_mqa(hidden_states) q, compressed_kv, k_pe = kv_a.split( [mla.q_lora_rank, mla.kv_lora_rank, mla.qk_rope_head_dim], -1 ) @@ -335,11 +474,10 @@ def _attention(layer, position_ids, hidden_states, attn_metadata, all_reduce_par indexer.maybe_join_prev_topk_copy() return mla.o_proj(output, all_reduce_params=all_reduce_params, layer_idx=mla.layer_idx) - @staticmethod def _dense_layer( - layer, position_ids, hidden_states, residual, attn_metadata, fused_a, small_m + self, layer, position_ids, hidden_states, residual, attn_metadata, fused_a, small_m ): - hidden_states = DecodeTarget._attention( + hidden_states = self._attention( layer, position_ids, hidden_states, @@ -368,7 +506,7 @@ def _dense_layer( ) def _moe_layer(self, layer, position_ids, hidden_states, residual, attn_metadata, fused_a): - hidden_states = DecodeTarget._attention( + hidden_states = self._attention( layer, position_ids, hidden_states, From d21070fa1caeabcf1cf84ebf02de2f359c53f2a8 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:09:36 -0700 Subject: [PATCH 22/53] Add the GLM-5.2 low-M BF16 MoE op for the one-token MTP experts torch.ops.trtllm.glm_lowm_bf16_moe runs one token through an MoE layer with unquantized BF16 SwiGLU experts held as tensor-parallel slices and one-group DeepSeek-V3 routing, in two kernels (glmMoeDsa/glmLowMBf16Moe.cu): - FC1 + SwiGLU, one CTA per SM: a producer warp routes from the router logits (noaux_tc_op's experts, order and tie-break) and streams whole weight rows of the routed experts through a shared-memory ring that eight consumer warps read two rows at a time; it releases its dependents right after its grid-dependency wait. - FC2 + the weighted expert sum: routes again before its own wait from bulk-copied logits and bias, stages each expert's weight slice in shared memory beside the FC1 kernel, forms noaux_tc_op's routing weights in FP32 and falls back to its FP64 expression only where the FP32 result is not certain. The BF16 rounding points are those of the CUTLASS MoE path (FC1 output, SwiGLU output, each expert's scaled FC2 output); the expert sum is FP32 in the routing order. 256 experts, top 8, hidden sizes that are multiples of 2048 up to 8192 and expert slices that are multiples of 64 up to 512; anything else, and more than one token, is refused. The tp16 GLM target's owned entry ops/moe/lowm_bf16_moe certifies the MTP layer's draft-step shape (hidden 6144, slice 128) over random, tied and tiny-score routing on sm107, plus a CUDA-graph replay and the refusals. Signed-off-by: Dongfeng Yu --- .../kernels/glmMoeDsa/glmLowMBf16Moe.cu | 839 ++++++++++++++++++ .../kernels/glmMoeDsa/glmLowMBf16Moe.h | 74 ++ cpp/tensorrt_llm/thop/CMakeLists.txt | 3 +- cpp/tensorrt_llm/thop/glmLowMBf16MoeOp.cpp | 124 +++ .../ops/moe/lowm_bf16_moe.py | 130 +++ .../_torch/custom_ops/cpp_custom_ops.py | 5 + .../moe/test_modeling_v2_lowm_bf16_moe.py | 147 +++ 7 files changed, 1321 insertions(+), 1 deletion(-) create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.cu create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.h create mode 100644 cpp/tensorrt_llm/thop/glmLowMBf16MoeOp.cpp create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/lowm_bf16_moe.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_lowm_bf16_moe.py diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.cu new file mode 100644 index 000000000000..850a0c2bcd7f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.cu @@ -0,0 +1,839 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/kernels/moeTopKFuncs.cuh" + +#include +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ +namespace +{ +namespace cg = cooperative_groups; + +using Bf16 = __nv_bfloat16; + +constexpr int kNumExperts = kLowMBf16MoeNumExperts; +constexpr int kTopK = kLowMBf16MoeTopK; +static_assert(kNumExperts % 32 == 0, "routing gives each lane of one warp kNumExperts / 32 candidates"); + +// FC1: 8 consumer warps and one producer warp per CTA. Consumer thread t reads the 16-byte pieces t, +// t + kFc1ConsumerThreads, ... of x and of every weight row, so the hidden size is a multiple of kHiddenGranule and at +// most kMaxPiecesPerThread of them. +constexpr int kFc1ConsumerWarps = 8; +constexpr int kFc1ConsumerThreads = kFc1ConsumerWarps * 32; +constexpr int kFc1Threads = kFc1ConsumerThreads + 32; +constexpr int kPieceElems = 8; +constexpr int kHiddenGranule = kFc1ConsumerThreads * kPieceElems; +constexpr int kMaxPiecesPerThread = 4; +// The FC1 ring holds whole weight rows, as many as fit in kFc1RingBytes and at most kMaxStages. At hidden size 6144 +// and interSize 256 that is 7 rows of 12,288 B, and an FC1 CTA (ring, row sums, barriers and the 1,024 B the runtime +// reserves per CTA) takes 87,840 B, so one FC2 CTA (138,312 B) fits beside it in the 233,472 B of an SM. +constexpr int kFc1RingBytes = 7 * 6144 * static_cast(sizeof(Bf16)); +constexpr int kMaxStages = 8; + +// FC2: a warp per routed expert. +constexpr int kFc2Threads = 256; +constexpr int kFc2Warps = kFc2Threads / 32; +static_assert(kTopK == kFc2Warps, "the FC2 kernel gives each routed expert one warp"); + +// FC2: mma.sync m16n8k16 tiles have 16 matrix rows; one 32-element K step is two mma calls (see mmaStep). A warp per +// routed expert walks the intermediate dimension in pairs of K steps, one per accumulator. +constexpr int kMmaRows = 16; +constexpr int kStepK = 32; +constexpr int kFc2StepPairK = 2 * kStepK; +// The largest intermediate size per rank whose FC2 working set fits in shared memory (kTiles = 1). +constexpr int kMaxInterSize = 512; + +// The routed scaling factor as the FC2 kernel takes it: the double value of noaux_tc_op's FP64 weight expression, its +// float value, and whether routingWeight's FP32 path applies to it (decided on the host, see makeRoutingScale). +struct RoutingScale +{ + double value; + float valueFloat; + int fastPath; +}; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +constexpr int kCandidates = kNumExperts / 32; +constexpr int kFc1ProducerWarp = kFc1ConsumerWarps; +// Named barrier at which the FC1 consumer warps wait while the producer selects the experts and issues the first rows. +constexpr int kFc1ParkBarrier = 1; +// The FC2 kernel holds its routing computation until this long after its entry; see lowMBf16MoeFc2Kernel. +constexpr uint64_t kFc2RouteHoldNs = 1000; + +using TopKKey = reduce_topk::TopKRedType; + +// Same expression as sigmoid_accurate in noAuxTcKernels.cu. +__device__ __forceinline__ float sigmoidAccurate(float x) +{ + return 0.5f * tanhf(0.5f * x) + 0.5f; +} + +__device__ __forceinline__ uint32_t warpMax(uint32_t value) +{ + uint32_t result; + asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n" : "=r"(result) : "r"(value)); + return result; +} + +// The correction bias of the experts a lane holds (i * 32 + lane) as float. A constant, so the kernels may read it +// before their grid-dependency wait. +template +__device__ __forceinline__ void loadBias(BiasT const* bias, int lane, float (&values)[kCandidates]) +{ +#pragma unroll + for (int i = 0; i < kCandidates; ++i) + { + values[i] = static_cast(bias[i * 32 + lane]); + } +} + +// The routing keys of noaux_tc_op for n_group == 1 (deepseek_v3_topk_kernel, single group) of the experts a lane +// holds: the sigmoid score plus the correction bias, with the expressions of that kernel, packed with the expert index +// as reduceTopK (moeTopKFuncs.cuh) packs them. Stores the sigmoid scores when sigmoidScores is not null. The logits +// are read through L2. +__device__ __forceinline__ void buildKeys(float const* routerLogits, float const (&biasValues)[kCandidates], int lane, + TopKKey (&keys)[kCandidates], float* sigmoidScores) +{ +#pragma unroll + for (int i = 0; i < kCandidates; ++i) + { + int const expert = i * 32 + lane; + float const sigmoid = sigmoidAccurate(__ldcg(routerLogits + expert)); + if (sigmoidScores != nullptr) + { + sigmoidScores[expert] = sigmoid; + } + keys[i] = TopKKey{sigmoid + biasValues[i], expert}; + } +} + +// The top kTopK of one full warp's keys in the order of reduceTopK: descending score, on equal scores the lower index +// first. The per-lane sort is reduceTopK's; each round then takes the warp maximum of the lanes' leading keys with two +// 32-bit reductions, of the score half and then of the index half among the lanes holding the maximum score, which is +// the maximum of reduceTopK's 64-bit reduction. So the experts and their order are noaux_tc_op's. Returns the k-th +// expert on lane k < kTopK, 0 on the other lanes. +__device__ __forceinline__ int32_t selectTopK(TopKKey (&keys)[kCandidates], int lane) +{ + reduce_topk::Sort::run(keys); + + int32_t laneExpert = 0; + uint64_t packedMax = 0; +#pragma unroll + for (int k = 0; k < kTopK; ++k) + { + if (k > 0 && packedMax == keys[0].compValIdx) + { + // This lane held the previous maximum: drop it, as reduceTopK does. +#pragma unroll + for (int i = 0; i + 1 < kCandidates; ++i) + { + keys[i] = keys[i + 1]; + } + keys[kCandidates - 1] = TopKKey{-INFINITY, (kCandidates - 1) * 32 + lane}; + } + uint32_t const scoreBits = static_cast(keys[0].compValIdx >> 32); + uint32_t const maxScoreBits = warpMax(scoreBits); + uint32_t const indexBits = warpMax(scoreBits == maxScoreBits ? static_cast(keys[0].compValIdx) : 0U); + packedMax = (static_cast(maxScoreBits) << 32) | indexBits; + if (lane == k) + { + laneExpert = TopKKey::kMaxIdx - static_cast(indexBits & 0xFFFF); + } + } + return laneExpert; +} + +// noaux_tc_op's routing weight, static_cast(s * c / (S + 1e-20)) in double with s the expert's sigmoid score, +// S the sum of the selected scores and c the routed scaling factor, without FP64 instructions on the common path: FP64 +// in the FC2 CTAs drops the SM clock for the whole FC1 window (measured on Rubin: 2.38 -> 1.6 GHz). When c is a float +// with at most 20 significant bits (fastPath) and S >= 2^-13, the double expression is the float rounding of the exact +// quotient s * c / S: s * c is exact in double, S + 1e-20 rounds to S, and rounding to double first cannot move a +// quotient of these operands across a float rounding boundary except at an exact midpoint, which rounds the same +// either way. The quotient here comes from an exact two-float product, a correctly rounded division and its exact +// remainder; it is returned when it lies more than 2^-20 ulp inside its rounding interval (and is not a power of two, +// where the interval is lopsided). Otherwise, and for scores below 2^-60 (where the FP32 intermediates could lose +// bits), the FP64 expression runs, out of line so that none of it is scheduled on the common path. +__device__ __noinline__ float routingWeightFp64(float s, float S, double c) +{ + return static_cast(s * c / (S + 1e-20)); +} + +__device__ __forceinline__ float routingWeight(float s, float S, RoutingScale const& scale) +{ + if (scale.fastPath != 0 && S >= 0x1p-13f && s >= 0x1p-60f) + { + float const c = scale.valueFloat; + float const pHi = __fmul_rn(s, c); + float const pLo = __fmaf_rn(s, c, -pHi); + float const q = __fdiv_rn(pHi, S); + float const r = __fmaf_rn(-q, S, pHi); + float const e = __fdiv_rn(__fadd_rn(r, pLo), S); + float const y = __fadd_rn(q, e); + uint32_t const bits = __float_as_uint(y); + if (y > 0.F && (bits & 0x7FFFFFU) != 0 && (bits & 0x7F800000U) != 0x7F800000U) + { + float const ulp = __uint_as_float(bits & 0x7F800000U) * 0x1p-23f; + float const g = __fsub_rn(e, __fsub_rn(y, q)); + if (fabsf(g) < 0.5F * ulp - 0x1p-20f * ulp) + { + return y; + } + } + } + return routingWeightFp64(s, S, scale.value); +} + +__device__ __forceinline__ uint32_t smemAddress(void const* ptr) +{ + return static_cast(__cvta_generic_to_shared(ptr)); +} + +__device__ __forceinline__ void mbarrierInit(uint64_t* bar, uint32_t count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(smemAddress(bar)), "r"(count) : "memory"); +} + +// Makes the barrier initialization visible to the bulk-copy (async) proxy. +__device__ __forceinline__ void fenceBarrierInit() +{ + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); +} + +__device__ __forceinline__ void mbarrierArriveExpectTx(uint64_t* bar, uint32_t bytes) +{ + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(smemAddress(bar)), "r"(bytes) + : "memory"); +} + +// Arrives on the barrier with release semantics: this thread's earlier shared-memory reads are done before the phase +// can complete. +__device__ __forceinline__ void mbarrierArrive(uint64_t* bar) +{ + asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];\n" ::"r"(smemAddress(bar)) : "memory"); +} + +// Whether the barrier's phase with the given parity has completed (one try, which may suspend the thread for a +// while). +__device__ __forceinline__ bool mbarrierTryWait(uint64_t* bar, uint32_t parity) +{ + uint32_t done = 0; + asm volatile( + "{\n" + ".reg .pred complete;\n" + "mbarrier.try_wait.parity.shared::cta.b64 complete, [%1], %2;\n" + "selp.u32 %0, 1, 0, complete;\n" + "}\n" + : "=r"(done) + : "r"(smemAddress(bar)), "r"(parity) + : "memory"); + return done != 0; +} + +__device__ __forceinline__ void mbarrierWait(uint64_t* bar, uint32_t parity) +{ + while (!mbarrierTryWait(bar, parity)) + { + } +} + +// Waits at named barrier `id` until `threads` threads (whole warps) have arrived. +__device__ __forceinline__ void namedBarrierSync(int id, int threads) +{ + asm volatile("bar.sync %0, %1;\n" ::"r"(id), "r"(threads) : "memory"); +} + +__device__ __forceinline__ uint64_t globalTimerNs() +{ + uint64_t now; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(now)); + return now; +} + +// Streamed weights are read once per call: mark their L2 lines evict-first. +__device__ __forceinline__ uint64_t evictFirstPolicy() +{ + uint64_t policy; + asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;\n" : "=l"(policy)); + return policy; +} + +// 1D bulk copy global -> shared that completes `bytes` transactions on `bar`. +__device__ __forceinline__ void bulkCopyToShared( + void* dst, void const* src, uint32_t bytes, uint64_t* bar, uint64_t policy) +{ + uint64_t const srcGlobal = static_cast(__cvta_generic_to_global(src)); + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], %4;\n" + : + : "r"(smemAddress(dst)), "l"(srcGlobal), "r"(bytes), "r"(smemAddress(bar)), "l"(policy) + : "memory"); +} + +__device__ __forceinline__ void mmaBf16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) +{ + asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +// One 32-element K step of a 16-row tile times a vector. Lane (g = lane / 4, t = lane % 4) passes elements [8t, 8t + 8) +// of the step from tile row g (lo), tile row g + 8 (hi) and the vector (vec). The fragment's K columns +// {2t, 2t + 1, 2t + 8, 2t + 9} take the lane's elements {0, 1, 2, 3} in the first call and {4, 5, 6, 7} in the second: +// one K permutation for the matrix and the vector, which leaves the dot products unchanged. Every column of the B +// fragment is the vector, so all accumulator elements hold dot products: elements 0 and 1 of row g, 2 and 3 of row +// g + 8. +__device__ __forceinline__ void mmaStep(float (&acc)[4], uint4 const& lo, uint4 const& hi, uint4 const& vec) +{ + mmaBf16(acc, lo.x, hi.x, lo.y, hi.y, vec.x, vec.y); + mmaBf16(acc, lo.z, hi.z, lo.w, hi.w, vec.z, vec.w); +} + +__device__ __forceinline__ float roundToBf16(float value) +{ + return __bfloat162float(__float2bfloat16_rn(value)); +} + +// The FP32 values of the eight BF16 elements of a 16-byte piece, in memory order. +__device__ __forceinline__ void unpackPiece(uint4 const& piece, float (&values)[kPieceElems]) +{ + uint32_t const words[4] = {piece.x, piece.y, piece.z, piece.w}; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + values[2 * i] = __uint_as_float(words[i] << 16); + values[2 * i + 1] = __uint_as_float(words[i] & 0xFFFF0000U); + } +} + +// Dot product of a 16-byte piece of a BF16 weight row with the FP32 values of the same piece of x, in element order +// with FP32 fused multiply-adds. +__device__ __forceinline__ float dotPiece(uint4 const& piece, float const (&values)[kPieceElems]) +{ + float weights[kPieceElems]; + unpackPiece(piece, weights); + float dot = weights[0] * values[0]; +#pragma unroll + for (int e = 1; e < kPieceElems; ++e) + { + dot = fmaf(weights[e], values[e], dot); + } + return dot; +} + +// The warp sums of two values, each by an xor butterfly over offsets 16, 8, 4, 2, 1, the two shuffle chains +// interleaved. +__device__ __forceinline__ void warpSum2(float& a, float& b) +{ +#pragma unroll + for (int offset = 16; offset > 0; offset /= 2) + { + a += __shfl_xor_sync(0xffffffffU, a, offset); + b += __shfl_xor_sync(0xffffffffU, b, offset); + } +} +#endif + +// FC1 + SwiGLU, one CTA per SM (the grid is the SM count). The work is kTopK * interSize units; unit k * interSize + j +// is intermediate index j of routed slot k, that is the expert's gate row (W1, w3w1 row interSize + j) and up row (W3, +// row j). CTA b takes units [b * U / G, (b + 1) * U / G) and streams their rows, gate then up, each as one bulk copy +// into a ring of numStages row buffers. The producer warp selects the experts and issues the rows; it refills a buffer +// as soon as every consumer warp has read its pieces of it, with no CTA-wide barrier per row. The consumer warps wait +// at a named barrier while the producer selects the experts and issues the first rows. Each consumer thread keeps its +// pieces of x in registers and adds its shares of the dot products of two rows at a time, so that the latency chains of +// the two overlap; each consumer warp reduces its shares with shuffles into shared memory, and each unit's SwiGLU adds +// the 8 warp sums in a fixed order. +template +__global__ void __launch_bounds__(kFc1Threads, 2) lowMBf16MoeFc1Kernel(Bf16* __restrict__ act, + Bf16 const* __restrict__ x, float const* __restrict__ routerLogits, BiasT const* __restrict__ bias, + Bf16 const* __restrict__ w3w1, int hiddenSize, int interSize, int numStages) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + extern __shared__ __align__(128) uint8_t smem[]; + // rowFull[s]: buffer s holds its next row (the producer's arrival plus the copy's bytes). rowFree[s]: every + // consumer warp has read its pieces of buffer s. + __shared__ uint64_t rowFull[kMaxStages]; + __shared__ uint64_t rowFree[kMaxStages]; + __shared__ int32_t topExperts[kTopK]; + + int const warp = threadIdx.x / 32; + int const lane = threadIdx.x % 32; + int const rowBytes = hiddenSize * static_cast(sizeof(Bf16)); + int const numUnits = kTopK * interSize; + int const unitBegin = static_cast(static_cast(blockIdx.x) * numUnits / gridDim.x); + int const unitEnd = static_cast(static_cast(blockIdx.x + 1) * numUnits / gridDim.x); + int const numRows = 2 * (unitEnd - unitBegin); + uint8_t* ring = smem; + float* rowSums = reinterpret_cast(smem + numStages * rowBytes); + + if (threadIdx.x == 0) + { + for (int s = 0; s < numStages; ++s) + { + mbarrierInit(&rowFull[s], 1); + mbarrierInit(&rowFree[s], kFc1ConsumerWarps); + } + fenceBarrierInit(); + } + // The correction bias is a constant: the producer reads it before the grid-dependency wait. + float biasValues[kCandidates] = {}; + if (warp == kFc1ProducerWarp) + { + loadBias(bias, lane, biasValues); + } + __syncthreads(); + + // x and the router logits are written by earlier kernels. + cudaGridDependencySynchronize(); + // The FC2 kernel reads only the logits, the bias and its weights before its own wait, and the logits are complete + // since the wait above returned. Released here, its CTAs take the room beside this kernel's and stage their + // weights while this kernel streams. + cudaTriggerProgrammaticLaunchCompletion(); + + if (warp == kFc1ProducerWarp) + { + TopKKey keys[kCandidates]; + buildKeys(routerLogits, biasValues, lane, keys, nullptr); + int32_t const expert = selectTopK(keys, lane); + if (lane < kTopK) + { + topExperts[lane] = expert; + } + __syncwarp(); + + // Lane 0 streams the rows. The next row is of unit (slot, index), the gate row first, into buffer s, which it + // fills for the pass-th time. + uint64_t const policy = evictFirstPolicy(); + int64_t const expertStride = 2 * static_cast(interSize) * hiddenSize; + int slot = unitBegin / interSize; + int index = unitBegin % interSize; + bool gateRow = true; + Bf16 const* expertRows = w3w1 + topExperts[slot] * expertStride; + int s = 0; + uint32_t pass = 0; + auto issueRow = [&]() + { + int64_t const rowInExpert = gateRow ? interSize + index : index; + mbarrierArriveExpectTx(&rowFull[s], rowBytes); + bulkCopyToShared(ring + s * rowBytes, expertRows + rowInExpert * hiddenSize, rowBytes, &rowFull[s], policy); + if (!gateRow && ++index == interSize) + { + // The up row closed the slot's last unit. + index = 0; + if (++slot < kTopK) + { + expertRows = w3w1 + topExperts[slot] * expertStride; + } + } + gateRow = !gateRow; + if (++s == numStages) + { + s = 0; + ++pass; + } + }; + int const firstRows = numStages < numRows ? numStages : numRows; + if (lane == 0) + { + for (int row = 0; row < firstRows; ++row) + { + issueRow(); + } + } + // Release the consumer warps. + __syncwarp(); + namedBarrierSync(kFc1ParkBarrier, kFc1Threads); + if (lane == 0) + { + for (int row = firstRows; row < numRows; ++row) + { + // Buffer s still holds row - numStages until every consumer warp has read it. + mbarrierWait(&rowFree[s], (pass - 1) & 1U); + issueRow(); + } + } + } + else + { + float xValues[kMaxPiecesPerThread][kPieceElems] = {}; + int const piecesPerThread = hiddenSize / kHiddenGranule; +#pragma unroll + for (int i = 0; i < kMaxPiecesPerThread; ++i) + { + if (i < piecesPerThread) + { + unpackPiece(__ldcg(reinterpret_cast(x) + threadIdx.x + i * kFc1ConsumerThreads), + xValues[i]); + } + } + // Wait at the barrier, without polling the ring, while the producer selects the experts and issues the first + // rows. + namedBarrierSync(kFc1ParkBarrier, kFc1Threads); + int s0 = 0; + uint32_t pass0 = 0; + // Rows row and row + 1 (numRows is even) are in buffers s0 and s1 (numStages >= 2), filled for the pass0-th and + // pass1-th time. + for (int row = 0; row < numRows; row += 2) + { + int s1 = s0 + 1; + uint32_t pass1 = pass0; + if (s1 == numStages) + { + s1 = 0; + ++pass1; + } + mbarrierWait(&rowFull[s0], pass0 & 1U); + mbarrierWait(&rowFull[s1], pass1 & 1U); + uint4 const* pieces0 = reinterpret_cast(ring + s0 * rowBytes); + uint4 const* pieces1 = reinterpret_cast(ring + s1 * rowBytes); + float dot0 = dotPiece(pieces0[threadIdx.x], xValues[0]); + float dot1 = dotPiece(pieces1[threadIdx.x], xValues[0]); +#pragma unroll + for (int i = 1; i < kMaxPiecesPerThread; ++i) + { + if (i < piecesPerThread) + { + dot0 += dotPiece(pieces0[threadIdx.x + i * kFc1ConsumerThreads], xValues[i]); + dot1 += dotPiece(pieces1[threadIdx.x + i * kFc1ConsumerThreads], xValues[i]); + } + } + // This warp is done with both buffers (the arrival releases its lanes' reads, ordered before it by the + // warp barrier); the producer refills a buffer once every consumer warp is. + __syncwarp(); + if (lane == 0) + { + mbarrierArrive(&rowFree[s0]); + mbarrierArrive(&rowFree[s1]); + } + warpSum2(dot0, dot1); + if (lane == 0) + { + rowSums[row * kFc1ConsumerWarps + warp] = dot0; + rowSums[(row + 1) * kFc1ConsumerWarps + warp] = dot1; + } + s0 = s1 + 1; + pass0 = pass1; + if (s0 == numStages) + { + s0 = 0; + ++pass0; + } + } + } + __syncthreads(); + + for (int unit = static_cast(threadIdx.x); unit < numRows / 2; unit += kFc1Threads) + { + float gate = 0.F; + float up = 0.F; + for (int w = 0; w < kFc1ConsumerWarps; ++w) + { + gate += rowSums[2 * unit * kFc1ConsumerWarps + w]; + up += rowSums[(2 * unit + 1) * kFc1ConsumerWarps + w]; + } + // The CUTLASS FC1 GEMM writes its output in BF16 and the activation kernel reads it back. + gate = roundToBf16(gate); + up = roundToBf16(up); + // silu(gate) * up in FP32 as the CUTLASS activation (GLUAdaptor) computes it, then BF16. + float const silu = gate * (1.F / (1.F + expf(-gate))); + act[unitBegin + unit] = __float2bfloat16_rn(silu * up); + } +#endif +} + +// FC2 + the weighted sum over the routed experts. Grid: one CTA per kTiles * 16 output rows; warp k computes routed +// expert k's terms from its weight slice, staged in shared memory before the grid-dependency wait. Warp 0 routes before +// that wait, from the router logits and the bias staged by bulk copies, without FP64 instructions. After the wait every +// warp stages its expert's SwiGLU outputs itself. +template +__global__ void __launch_bounds__(kFc2Threads, 2) lowMBf16MoeFc2Kernel(Bf16* __restrict__ out, + Bf16 const* __restrict__ act, float const* __restrict__ routerLogits, BiasT const* __restrict__ bias, + Bf16 const* __restrict__ w2, int hiddenSize, int interSize, RoutingScale scale) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + constexpr int kRows = kTiles * kMmaRows; + extern __shared__ __align__(128) uint8_t smem[]; + __shared__ __align__(16) float logitsShared[kNumExperts]; + __shared__ __align__(16) BiasT biasShared[kNumExperts]; + __shared__ uint64_t inputsFull; + __shared__ uint64_t slicesFull; + __shared__ float sigmoidScores[kNumExperts]; + __shared__ int32_t topExperts[kTopK]; + __shared__ float topWeights[kTopK]; + __shared__ float terms[kTopK][kRows]; + + uint64_t const entryNs = globalTimerNs(); + + int const sliceElems = kRows * interSize; + Bf16* slices = reinterpret_cast(smem); + Bf16* actShared = slices + kTopK * sliceElems; + + int const warp = threadIdx.x / 32; + int const lane = threadIdx.x % 32; + int const row0 = blockIdx.x * kRows; + + if (threadIdx.x == 0) + { + mbarrierInit(&inputsFull, 1); + mbarrierInit(&slicesFull, 1); + fenceBarrierInit(); + } + __syncthreads(); + + if (warp == 0) + { + // Before the grid-dependency wait this kernel reads only the router logits, the bias and the FC2 weights. The + // logits come from the router GEMM, which completed before the FC1 kernel returned from its own wait and + // released this grid; the bias and the weights are constants. The logits and the bias come in by bulk copies: + // plain loads issued now queue behind the FC1 kernel's bulk copies on this SM for microseconds (measured). + bool const bulkInputs + = ((reinterpret_cast(routerLogits) | reinterpret_cast(bias)) % 16) == 0; + if (bulkInputs) + { + if (lane == 0) + { + uint64_t const policy = evictFirstPolicy(); + mbarrierArriveExpectTx(&inputsFull, sizeof(logitsShared) + sizeof(biasShared)); + bulkCopyToShared(logitsShared, routerLogits, sizeof(logitsShared), &inputsFull, policy); + bulkCopyToShared(biasShared, bias, sizeof(biasShared), &inputsFull, policy); + } + mbarrierWait(&inputsFull, 0); + } + else + { + // Each lane stages, and below reads back, only the experts it holds. +#pragma unroll + for (int i = 0; i < kCandidates; ++i) + { + logitsShared[i * 32 + lane] = __ldcg(routerLogits + i * 32 + lane); + biasShared[i * 32 + lane] = bias[i * 32 + lane]; + } + } + // The FC1 kernel's producer warps select the same experts on these SMs at about this time, on the FC1 + // kernel's critical path, and routing here concurrently makes their selection about twice as long (measured + // 0.5 -> 1.1 us). So the routing waits until kFc2RouteHoldNs after this CTA's entry; the weight copies are + // still issued long before the FC1 kernel ends. + while (globalTimerNs() < entryNs + kFc2RouteHoldNs) + { + __nanosleep(64); + } + // The FC1 kernel's routing (the keys of buildKeys, selectTopK) on the staged inputs, and the weights from + // routingWeight: the experts, their order and the FP32 weights of noaux_tc_op. + TopKKey keys[kCandidates]; +#pragma unroll + for (int i = 0; i < kCandidates; ++i) + { + int const expert = i * 32 + lane; + float const sigmoid = sigmoidAccurate(logitsShared[expert]); + sigmoidScores[expert] = sigmoid; + keys[i] = TopKKey{sigmoid + static_cast(biasShared[expert]), expert}; + } + int32_t const expert = selectTopK(keys, lane); + __syncwarp(); + auto const tile = cg::tiled_partition<32>(cg::this_thread_block()); + float const scoreNorm = lane < kTopK ? sigmoidScores[expert] : 0.F; + float const sumNorm = cg::reduce(tile, scoreNorm, cg::plus{}); + if (lane < kTopK) + { + topExperts[lane] = expert; + topWeights[lane] = routingWeight(scoreNorm, sumNorm, scale); + } + __syncwarp(); + if (lane == 0) + { + mbarrierArriveExpectTx(&slicesFull, kTopK * sliceElems * static_cast(sizeof(Bf16))); + } + __syncwarp(); + if (lane < kTopK) + { + bulkCopyToShared(slices + lane * sliceElems, + w2 + (static_cast(topExperts[lane]) * hiddenSize + row0) * interSize, + sliceElems * static_cast(sizeof(Bf16)), &slicesFull, evictFirstPolicy()); + } + } + __syncthreads(); + + // The SwiGLU outputs are written by the FC1 kernel: every warp waits for it, then stages its expert's outputs. + cudaGridDependencySynchronize(); + Bf16 const* slice = slices + warp * sliceElems; + Bf16* expertAct = actShared + warp * interSize; + for (int i = lane; i < interSize / 8; i += 32) + { + reinterpret_cast(expertAct)[i] = __ldcg(reinterpret_cast(act + warp * interSize) + i); + } + __syncwarp(); + mbarrierWait(&slicesFull, 0); + + // Lane (g, t) of warp k accumulates the dot products of slice rows tile * 16 + g and tile * 16 + g + 8 with + // act[k, :], over the two 32-element halves of each 64-element K step in two accumulators. Lanes of odd g take the + // two halves in the other order, so that a warp's loads of 8 rows spread over all 32 banks: column n of the B + // fragments then holds the vector in the order of the rows of n's parity, and the accumulators of those columns + // hold the same products in the same order as with one order for all (which of the two holds which half differs). + int const g = lane / 4; + int const t = lane % 4; + int const swap = g & 1; + float acc[kTiles][2][4] = {}; + for (int j0 = 0; j0 < interSize; j0 += kFc2StepPairK) + { +#pragma unroll + for (int half = 0; half < 2; ++half) + { + int const j = j0 + (half ^ swap) * kStepK + 8 * t; + uint4 const av = *reinterpret_cast(expertAct + j); +#pragma unroll + for (int tileIdx = 0; tileIdx < kTiles; ++tileIdx) + { + uint4 const lo = *reinterpret_cast(slice + (tileIdx * kMmaRows + g) * interSize + j); + uint4 const hi + = *reinterpret_cast(slice + (tileIdx * kMmaRows + g + 8) * interSize + j); + mmaStep(acc[tileIdx][half], lo, hi, av); + } + } + } + + // The CUTLASS FC2 epilogue rounds each expert's scaled output to BF16 before adding it into the output. Lane (g, 0) + // holds columns 0 and 1 of rows g and g + 8, of which column g & 1 has their dot products. + float const weight = topWeights[warp]; + if (t == 0) + { +#pragma unroll + for (int tileIdx = 0; tileIdx < kTiles; ++tileIdx) + { + float const lo0 = swap ? acc[tileIdx][0][1] : acc[tileIdx][0][0]; + float const lo1 = swap ? acc[tileIdx][1][1] : acc[tileIdx][1][0]; + float const hi0 = swap ? acc[tileIdx][0][3] : acc[tileIdx][0][2]; + float const hi1 = swap ? acc[tileIdx][1][3] : acc[tileIdx][1][2]; + terms[warp][tileIdx * kMmaRows + g] = roundToBf16(weight * (lo0 + lo1)); + terms[warp][tileIdx * kMmaRows + g + 8] = roundToBf16(weight * (hi0 + hi1)); + } + } + __syncthreads(); + if (threadIdx.x < kRows) + { + float sum = 0.F; +#pragma unroll + for (int k = 0; k < kTopK; ++k) + { + sum += terms[k][threadIdx.x]; + } + out[row0 + threadIdx.x] = __float2bfloat16_rn(sum); + } +#endif +} + +// Lets a kernel use smemBytes of dynamic shared memory and asks for the largest shared-memory carveout for both +// kernels, so that an SM running a CTA of one can take a CTA of the other. +template +void setSharedMemoryAttributes(Kernel* kernel, size_t smemBytes) +{ + TLLM_CUDA_CHECK( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smemBytes))); + TLLM_CUDA_CHECK(cudaFuncSetAttribute( + kernel, cudaFuncAttributePreferredSharedMemoryCarveout, static_cast(cudaSharedmemCarveoutMaxShared))); +} + +// routingWeight's FP32 path applies when the scaling factor is a float with at most 20 significant bits (the low 4 +// mantissa bits zero) between 2^-20 and 2^20. +RoutingScale makeRoutingScale(double routedScalingFactor) +{ + float const valueFloat = static_cast(routedScalingFactor); + uint32_t bits = 0; + std::memcpy(&bits, &valueFloat, sizeof(bits)); + bool const fastPath = static_cast(valueFloat) == routedScalingFactor && (bits & 0xFU) == 0 + && valueFloat >= 0x1p-20f && valueFloat <= 0x1p20f; + return RoutingScale{routedScalingFactor, valueFloat, fastPath ? 1 : 0}; +} + +template +void launchFc2(Bf16* out, Bf16 const* act, float const* routerLogits, BiasT const* bias, Bf16 const* w2, + int hiddenSize, int interSize, double routedScalingFactor, cudaStream_t stream) +{ + int const rows = kTiles * kMmaRows; + size_t const smemBytes = static_cast(kTopK) * (rows + 1) * interSize * sizeof(Bf16); + auto* kernel = &lowMBf16MoeFc2Kernel; + setSharedMemoryAttributes(kernel, smemBytes); + common::launchWithPdlWhenEnabled("lowMBf16MoeFc2", kernel, dim3(hiddenSize / rows), dim3(kFc2Threads), smemBytes, + stream, out, act, routerLogits, bias, w2, hiddenSize, interSize, makeRoutingScale(routedScalingFactor)); +} + +} // namespace + +bool lowMBf16MoeSupportsShape(int64_t hiddenSize, int64_t interSize) +{ + return hiddenSize > 0 && hiddenSize % kHiddenGranule == 0 && hiddenSize <= kMaxPiecesPerThread * kHiddenGranule + && interSize > 0 && interSize % kFc2StepPairK == 0 && interSize <= kMaxInterSize; +} + +template +void invokeLowMBf16Moe(Bf16* out, Bf16* act, Bf16 const* x, float const* routerLogits, BiasT const* bias, + Bf16 const* w3w1, Bf16 const* w2, int hiddenSize, int interSize, double routedScalingFactor, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(lowMBf16MoeSupportsShape(hiddenSize, interSize), + "lowMBf16Moe: unsupported hidden size %d or intermediate size %d", hiddenSize, interSize); + + // FC1: one CTA per SM, each with an even share of the units and a slot per row and consumer warp for the sums. + static int const smCount = common::getMultiProcessorCount(); + int const numUnits = kTopK * interSize; + int const fc1Grid = std::min(smCount, numUnits); + int const rowBytes = hiddenSize * static_cast(sizeof(Bf16)); + int const numStages = std::min(kMaxStages, kFc1RingBytes / rowBytes); + int const maxRows = 2 * ((numUnits + fc1Grid - 1) / fc1Grid); + size_t const fc1SmemBytes = static_cast(numStages) * rowBytes + + static_cast(maxRows) * kFc1ConsumerWarps * sizeof(float); + auto* fc1 = &lowMBf16MoeFc1Kernel; + setSharedMemoryAttributes(fc1, fc1SmemBytes); + common::launchWithPdlWhenEnabled("lowMBf16MoeFc1", fc1, dim3(fc1Grid), dim3(kFc1Threads), fc1SmemBytes, stream, + act, x, routerLogits, bias, w3w1, hiddenSize, interSize, numStages); + + // Two 16-row tiles per FC2 CTA where the 8 weight slices fit in shared memory, else one. + if (interSize <= kMaxInterSize / 2) + { + launchFc2(out, act, routerLogits, bias, w2, hiddenSize, interSize, routedScalingFactor, stream); + } + else + { + launchFc2(out, act, routerLogits, bias, w2, hiddenSize, interSize, routedScalingFactor, stream); + } +} + +#define INSTANTIATE_LOW_M_BF16_MOE(BiasT) \ + template void invokeLowMBf16Moe(Bf16 * out, Bf16 * act, Bf16 const* x, float const* routerLogits, \ + BiasT const* bias, Bf16 const* w3w1, Bf16 const* w2, int hiddenSize, int interSize, \ + double routedScalingFactor, cudaStream_t stream); + +INSTANTIATE_LOW_M_BF16_MOE(float); +INSTANTIATE_LOW_M_BF16_MOE(__nv_bfloat16); + +#undef INSTANTIATE_LOW_M_BF16_MOE + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.h new file mode 100644 index 000000000000..928eebda65e2 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.h @@ -0,0 +1,74 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +//! Routed experts and experts per token the low-M BF16 MoE kernels are built for. +constexpr int kLowMBf16MoeNumExperts = 256; +constexpr int kLowMBf16MoeTopK = 8; + +//! Whether invokeLowMBf16Moe supports a layer: hiddenSize a multiple of 2048 and at most 8192, interSize a multiple +//! of 64 and at most 512. +bool lowMBf16MoeSupportsShape(int64_t hiddenSize, int64_t interSize); + +//! One token through an MoE layer with unquantized BF16 SwiGLU experts held as tensor-parallel slices (no expert +//! parallelism) and DeepSeek-V3 routing with one expert group, in two kernels: +//! 1. route, then for each routed expert e_k: act[k, :] = bf16(silu(bf16(W1[e_k] x)) * bf16(W3[e_k] x)); +//! 2. route again, then out = bf16(sum_k bf16(w_k * (W2[e_k] act[k, :]))). +//! Routing is the noaux_tc_op computation for n_group == 1 (deepseek_v3_topk_kernel): the same experts, in the same +//! order, and bit-identical weights w_k (normalized sigmoid scores times routedScalingFactor). The BF16 rounding points +//! are those of the CUTLASS MoE path: the FC1 output, the SwiGLU output and each expert's scaled FC2 output. The sum +//! over the experts is in FP32 in a fixed order, where the CUTLASS finalize adds BF16 terms with atomics. +//! +//! The FC1 kernel runs one CTA per SM: a producer warp selects the experts and streams whole weight rows through a +//! shared-memory ring, which 8 consumer warps read two rows at a time. It releases the FC2 kernel right after its +//! grid-dependency wait, so that FC2 CTAs fit beside it and stage their weights in shared memory while FC1 streams. +//! Neither kernel runs FP64 instructions on its common path (FP64 lowers the SM clock): the FC2 kernel forms the +//! weights in FP32 with the result of noaux_tc_op's FP64 expression, which it runs only where that is not certain. +//! +//! Under PDL the FC1 kernel reads x and the logits after its grid-dependency wait; the FC2 kernel reads only the +//! logits, the bias and the FC2 weights before its own wait, which is safe because the FC1 kernel releases it after +//! its wait has returned. +//! +//! \param out Output [hiddenSize]. +//! \param act Scratch [kLowMBf16MoeTopK, interSize] for the SwiGLU outputs, 16-byte aligned. +//! \param x Input token [hiddenSize], 16-byte aligned. +//! \param routerLogits Router logits [kLowMBf16MoeNumExperts], FP32. +//! \param bias Routing correction bias [kLowMBf16MoeNumExperts]. +//! \param w3w1 FC1 weights [kLowMBf16MoeNumExperts, 2 * interSize, hiddenSize], 16-byte aligned: rows +//! [0, interSize) of an expert are W3 (up), rows [interSize, 2 * interSize) are W1 (gate), the +//! layout CutlassFusedMoE stores. +//! \param w2 FC2 weights [kLowMBf16MoeNumExperts, hiddenSize, interSize], 16-byte aligned. +template +void invokeLowMBf16Moe(__nv_bfloat16* out, __nv_bfloat16* act, __nv_bfloat16 const* x, float const* routerLogits, + BiasT const* bias, __nv_bfloat16 const* w3w1, __nv_bfloat16 const* w2, int hiddenSize, int interSize, + double routedScalingFactor, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 16b2633a737f..d50b691a6b46 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -122,7 +122,8 @@ add_th_common_family( # because moe/utils/moeUtilOp.cpp depends on OSS-specific MoE utility types and # function signatures. add_th_common_family(th_common_moe fusedTopkSoftmax.cpp llama4MinLatency.cpp - noAuxTcOp.cpp dsv3RouterGemmOp.cpp glmRouterGemmOp.cpp) + noAuxTcOp.cpp dsv3RouterGemmOp.cpp glmRouterGemmOp.cpp + glmLowMBf16MoeOp.cpp) # MoE torch ops live in moe//; that directory adds them to # th_common_moe with target_sources() so the list stays next to the files. diff --git a/cpp/tensorrt_llm/thop/glmLowMBf16MoeOp.cpp b/cpp/tensorrt_llm/thop/glmLowMBf16MoeOp.cpp new file mode 100644 index 000000000000..383d6c0cb676 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmLowMBf16MoeOp.cpp @@ -0,0 +1,124 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ +bool isAligned16(at::Tensor const& t) +{ + return reinterpret_cast(t.data_ptr()) % 16 == 0; +} +} // namespace + +// One token through an MoE layer with unquantized BF16 SwiGLU experts held as TP slices and one-group DeepSeek-V3 +// routing: routing from the router logits, FC1 + SwiGLU, and FC2 + the weighted sum over the routed experts, in two +// kernels (glmLowMBf16Moe.h). The weights are read in the layout CutlassFusedMoE keeps: w3_w1 [E, 2I, H] with the up +// (W3) rows first, w2 [E, H, I]. Returns [1, H]. +at::Tensor glmLowmBf16Moe(at::Tensor const& x, at::Tensor const& routerLogits, at::Tensor const& eScoreCorrectionBias, + at::Tensor const& w3W1, at::Tensor const& w2, int64_t topK, double routedScalingFactor) +{ + namespace tk = tensorrt_llm::kernels::glm_moe_dsa; + + CHECK_INPUT(x, at::ScalarType::BFloat16); + CHECK_INPUT(routerLogits, at::ScalarType::Float); + CHECK_TH_CUDA(eScoreCorrectionBias); + CHECK_CONTIGUOUS(eScoreCorrectionBias); + CHECK_INPUT(w3W1, at::ScalarType::BFloat16); + CHECK_INPUT(w2, at::ScalarType::BFloat16); + TORCH_CHECK(eScoreCorrectionBias.scalar_type() == at::ScalarType::Float + || eScoreCorrectionBias.scalar_type() == at::ScalarType::BFloat16, + "glm_lowm_bf16_moe expects an FP32 or BF16 e_score_correction_bias, got ", + eScoreCorrectionBias.scalar_type()); + for (at::Tensor const* t : {&routerLogits, &eScoreCorrectionBias, &w3W1, &w2}) + { + TORCH_CHECK(t->device() == x.device(), "glm_lowm_bf16_moe: all tensors must be on the device of x"); + } + TORCH_CHECK(tensorrt_llm::common::getSMVersion() >= 90, "glm_lowm_bf16_moe needs SM90 or newer"); + + TORCH_CHECK(x.dim() == 2 && x.size(0) == 1, "glm_lowm_bf16_moe takes one token, got x of shape ", x.sizes()); + TORCH_CHECK(w3W1.dim() == 3 && w2.dim() == 3, "glm_lowm_bf16_moe expects 3D expert weights"); + int64_t const hiddenSize = x.size(1); + int64_t const numExperts = w3W1.size(0); + int64_t const interSize = w2.size(2); + TORCH_CHECK(numExperts == tk::kLowMBf16MoeNumExperts, "glm_lowm_bf16_moe is built for ", + tk::kLowMBf16MoeNumExperts, " experts, got ", numExperts); + TORCH_CHECK(topK == tk::kLowMBf16MoeTopK, "glm_lowm_bf16_moe is built for top-", tk::kLowMBf16MoeTopK, + ", got top-", topK); + TORCH_CHECK(w3W1.size(1) == 2 * interSize && w3W1.size(2) == hiddenSize, "glm_lowm_bf16_moe: w3_w1 of shape ", + w3W1.sizes(), " does not match x ", x.sizes(), " and w2 ", w2.sizes()); + TORCH_CHECK(w2.size(0) == numExperts && w2.size(1) == hiddenSize, "glm_lowm_bf16_moe: w2 of shape ", w2.sizes(), + " does not match x ", x.sizes(), " and w3_w1 ", w3W1.sizes()); + TORCH_CHECK(routerLogits.dim() == 2 && routerLogits.size(0) == 1 && routerLogits.size(1) == numExperts, + "glm_lowm_bf16_moe: router_logits must be [1, ", numExperts, "], got ", routerLogits.sizes()); + TORCH_CHECK(eScoreCorrectionBias.dim() == 1 && eScoreCorrectionBias.size(0) == numExperts, + "glm_lowm_bf16_moe: e_score_correction_bias must be [", numExperts, "], got ", eScoreCorrectionBias.sizes()); + TORCH_CHECK(tk::lowMBf16MoeSupportsShape(hiddenSize, interSize), "glm_lowm_bf16_moe: unsupported hidden size ", + hiddenSize, " or intermediate size ", interSize); + TORCH_CHECK(isAligned16(x) && isAligned16(w3W1) && isAligned16(w2), + "glm_lowm_bf16_moe needs 16-byte aligned x, w3_w1 and w2"); + + at::Tensor out = at::empty({1, hiddenSize}, x.options()); + at::Tensor act = at::empty({topK, interSize}, x.options()); + auto const stream = at::cuda::getCurrentCUDAStream(x.get_device()); + auto* outPtr = static_cast<__nv_bfloat16*>(out.data_ptr()); + auto* actPtr = static_cast<__nv_bfloat16*>(act.data_ptr()); + auto const* xPtr = static_cast<__nv_bfloat16 const*>(x.data_ptr()); + auto const* logitsPtr = static_cast(routerLogits.data_ptr()); + auto const* w3W1Ptr = static_cast<__nv_bfloat16 const*>(w3W1.data_ptr()); + auto const* w2Ptr = static_cast<__nv_bfloat16 const*>(w2.data_ptr()); + if (eScoreCorrectionBias.scalar_type() == at::ScalarType::Float) + { + tk::invokeLowMBf16Moe(outPtr, actPtr, xPtr, logitsPtr, + static_cast(eScoreCorrectionBias.data_ptr()), w3W1Ptr, w2Ptr, static_cast(hiddenSize), + static_cast(interSize), routedScalingFactor, stream); + } + else + { + tk::invokeLowMBf16Moe<__nv_bfloat16>(outPtr, actPtr, xPtr, logitsPtr, + static_cast<__nv_bfloat16 const*>(eScoreCorrectionBias.data_ptr()), w3W1Ptr, w2Ptr, + static_cast(hiddenSize), static_cast(interSize), routedScalingFactor, stream); + } + return out; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_lowm_bf16_moe(Tensor x, Tensor router_logits, Tensor e_score_correction_bias, Tensor w3_w1, Tensor w2, " + "int top_k, float routed_scaling_factor) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_lowm_bf16_moe", &tensorrt_llm::torch_ext::glmLowmBf16Moe); +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/lowm_bf16_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/lowm_bf16_moe.py new file mode 100644 index 000000000000..f2b83f9f4ede --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/lowm_bf16_moe.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One token through the MTP layer's BF16 MoE, routing included, via the glm_lowm_bf16_moe op.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + + +class GlmLowmBf16Moe(OpWrapper): + """The routed half of the MTP layer's MoE for one token, from the router logits to the weighted expert sum. + + Routing is the n_group == 1 case of `noaux_tc_op`: experts by `sigmoid(logits) + bias` (sigmoid in the tanh + form), the top 8 in descending order with the lower expert id first on ties, weights from the unbiased scores + normalized over the 8 and scaled by `routed_scaling_factor`. Then per routed expert `e_k`: + `act_k = bf16(silu(bf16(W1[e_k] x)) * bf16(W3[e_k] x))` and the output is + `bf16(sum_k bf16(w_k * (W2[e_k] act_k)))`, the sum in FP32 in the routing order. + + `w3_w1` is `[256, 2 * I, H]` with the up rows (W3) first and the gate rows (W1) second, `w2` is `[256, H, I]`, + all experts local (tensor-parallel slices, no expert parallelism). Returns `[1, H]`; the shared experts are not + part of the op. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The MTP layer (layer 78) is BF16 throughout; at tp16 an expert's intermediate slice is 2048 / 16 = 128. Its + # MoE runs on one row on every draft step (step 0 after the row select). The checkpoint's correction bias is + # FP32 and its routed scaling factor 2.5. The value cells drive the routing's three regimes. + CELLS: tuple[Cell, ...] = ( + Cell( + why="draft step at tp16: one token, hidden 6144, expert slice 128, random routing", + spec=dict(hidden=6144, inter=128, routing="random"), + ), + Cell( + why="tied routing keys at the top-8 boundary: the lower expert id wins", + spec=dict(hidden=6144, inter=128, routing="ties"), + ), + Cell( + why="selected scores summing below 2^-13: the weights take the FP64 expression", + spec=dict(hidden=6144, inter=128, routing="tiny"), + ), + ) + + note = """ + The kernels serve hidden sizes that are multiples of 2048 up to 8192 and expert slices that are multiples of + 64 up to 512, exactly 256 experts and top 8; the op refuses anything else, and more than one token, loudly. A + BF16 correction bias is accepted as well; the checkpoint's is FP32, so no cell drives it. + + The order of the two halves of `w3_w1` is not visible in the shapes: a gate-first tensor gives silently wrong + output. The target builds `w3_w1` up-first from the checkpoint's `up_proj` and `gate_proj`. + + Under PDL the FC2 kernel reads the router logits before its grid-dependency wait: the logits must be complete + when the FC1 kernel's wait returns, which holds for any producer earlier in stream order. + """ + + def __call__( + self, + x: torch.Tensor, + router_logits: torch.Tensor, + e_score_correction_bias: torch.Tensor, + w3_w1: torch.Tensor, + w2: torch.Tensor, + top_k: int, + routed_scaling_factor: float, + ) -> torch.Tensor: + return torch.ops.trtllm.glm_lowm_bf16_moe( + x, router_logits, e_score_correction_bias, w3_w1, w2, top_k, routed_scaling_factor + ) + + def reference( + self, + x: torch.Tensor, + router_logits: torch.Tensor, + e_score_correction_bias: torch.Tensor, + w3_w1: torch.Tensor, + w2: torch.Tensor, + top_k: int, + routed_scaling_factor: float, + ) -> torch.Tensor: + """Plain torch with the op's rounding points. + + BF16 after FC1, after SwiGLU and after each expert's scaled FC2 output, the routing weights in FP64 as + `noaux_tc_op` forms them, the expert sum in FP32 in the routing order. Products and sums are FP32 (no + TF32). + """ + inter = w2.shape[2] + scores = 0.5 * torch.tanh(0.5 * router_logits.float()[0]) + 0.5 + choice = scores + e_score_correction_bias.float() + ids = torch.sort(choice, descending=True, stable=True).indices[:top_k] + selected = scores[ids] + total = selected.sum().double() + weights = (selected.double() * routed_scaling_factor / (total + 1e-20)).float() + xf = x.float()[0] + out = torch.zeros(w2.shape[1], dtype=torch.float32, device=x.device) + with true_fp32_matmul(): + for k in range(top_k): + expert = int(ids[k]) + fc1 = w3_w1[expert].float() @ xf + up = fc1[:inter].to(torch.bfloat16).float() + gate = fc1[inter:].to(torch.bfloat16).float() + act = (gate * (1.0 / (1.0 + torch.exp(-gate))) * up).to(torch.bfloat16).float() + term = (weights[k] * (w2[expert].float() @ act)).to(torch.bfloat16).float() + out = out + term + return out.to(torch.bfloat16).unsqueeze(0) + + def is_valid( + self, + x: torch.Tensor, + router_logits: torch.Tensor, + e_score_correction_bias: torch.Tensor, + w3_w1: torch.Tensor, + w2: torch.Tensor, + top_k: int, + routed_scaling_factor: float, + ) -> None: + # Nothing to add: the op refuses every token count, expert count, top-k, shape, dtype, layout and + # alignment it does not serve, loudly, and the one silent mistake left -- a gate-first `w3_w1` -- is not + # visible in the tensors (see `note`). + del x, router_logits, e_score_correction_bias, w3_w1, w2, top_k, routed_scaling_factor + + +glm_lowm_bf16_moe = GlmLowmBf16Moe() diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 6a998e291ae7..fea831399cb6 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -349,6 +349,11 @@ def _(scores, scores_with_bias, n_group, topk_group, topk, dtype=scores_with_bias.dtype), scores.new_empty( shape, dtype=torch.int32) + @torch.library.register_fake("trtllm::glm_lowm_bf16_moe") + def _(x, router_logits, e_score_correction_bias, w3_w1, w2, top_k, + routed_scaling_factor): + return x.new_empty((x.shape[0], w2.shape[1])) + @torch.library.register_fake("trtllm::kimi_k3_noaux_tc_mxfp8_quant") def _(scores, bias, hidden_states, routed_scaling_factor): num_tokens = scores.shape[0] diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_lowm_bf16_moe.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_lowm_bf16_moe.py new file mode 100644 index 000000000000..8bdec01cf856 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_lowm_bf16_moe.py @@ -0,0 +1,147 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm_lowm_bf16_moe entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.moe.lowm_bf16_moe import ( # noqa: E501 + glm_lowm_bf16_moe as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" +) + +NUM_EXPERTS = 256 +TOP_K = 8 +ROUTED_SCALING_FACTOR = 2.5 + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """One accumulation order's worth of BF16 rounding flips. + + The reference rounds at the op's points (FC1, SwiGLU, each scaled expert output, the sum), so the two differ + only where the FP32 accumulation order of a dot product moves a BF16 rounding of an intermediate. + + Over 6144-term FC1 dots of unit-scale operands the order changes the FP32 sum by a few 1e-6 relative, which + flips the BF16 rounding of about one gate or up value in a thousand; a flipped activation moves its expert's + FC2 outputs by about one BF16 ulp of the activation times a 1/sqrt(128)-scale weight, far below one ulp of the + output. What reaches the output is then the occasional flip of a term's or the sum's own rounding: one ulp of + that element, which is at most one ulp of the row's scale. + + So: 2 ulp element-wise against the row scale, 0.5 ulp of relative RMS. A wrong expert or weight moves the + output by about its own size, hundreds of ulp. + """ + assert_within_ulp(out, ref, element_ulp=2.0, rms_ulp=0.5) + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + hidden, inter = spec["hidden"], spec["inter"] + x = torch.randn(1, hidden, generator=g, device="cuda", dtype=torch.bfloat16) + w3_w1 = torch.randn( + NUM_EXPERTS, 2 * inter, hidden, generator=g, device="cuda", dtype=torch.bfloat16 + ).mul_(hidden**-0.5) + w2 = torch.randn( + NUM_EXPERTS, hidden, inter, generator=g, device="cuda", dtype=torch.bfloat16 + ).mul_(inter**-0.5) + routing = spec["routing"] + if routing == "random": + logits = 2.0 * torch.randn(1, NUM_EXPERTS, generator=g, device="cuda") + bias = 0.1 * torch.randn(NUM_EXPERTS, generator=g, device="cuda") + elif routing == "ties": + # Half-integer logits in [-2, 2] and no bias: about 28 experts share the top key, so which 8 of them are + # selected, and in which order, is the tie-break. + logits = ( + torch.randint(-4, 5, (1, NUM_EXPERTS), generator=g, device="cuda").float() / 2.0 + ) + bias = torch.zeros(NUM_EXPERTS, device="cuda") + else: + assert routing == "tiny" + # sigmoid(-16 +- 1) is about 1e-7, so the 8 selected scores sum to about 1e-6 < 2^-13. + logits = -16.0 + 0.5 * torch.randn(1, NUM_EXPERTS, generator=g, device="cuda") + bias = 0.1 * torch.randn(NUM_EXPERTS, generator=g, device="cuda") + return x, logits.contiguous(), bias.float().contiguous(), w3_w1, w2 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + out = op(*args, TOP_K, ROUTED_SCALING_FACTOR) + torch.cuda.synchronize() + _compare(out, op.reference(*args, TOP_K, ROUTED_SCALING_FACTOR)) + + +def test_cuda_graph_replay_matches_eager() -> None: + """The draft loop runs inside the decode graph: a replay with new inputs gives the eager result bitwise.""" + spec = op.CELLS[0].spec + x, logits, bias, w3_w1, w2 = _build(spec, seed=1) + static_x, static_logits = x.clone(), logits.clone() + out = op(static_x, static_logits, bias, w3_w1, w2, TOP_K, ROUTED_SCALING_FACTOR) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_out = op(static_x, static_logits, bias, w3_w1, w2, TOP_K, ROUTED_SCALING_FACTOR) + x2, logits2, _, _, _ = _build(spec, seed=2) + static_x.copy_(x2) + static_logits.copy_(logits2) + graph.replay() + eager = op(x2, logits2, bias, w3_w1, w2, TOP_K, ROUTED_SCALING_FACTOR) + torch.cuda.synchronize() + assert torch.equal(graph_out, eager) + assert not torch.equal(out, eager), "the replay must have read the new inputs" + + +def _small(hidden=2048, inter=64, tokens=1): + """Operands of the right layout at a small supported shape. + + The weights stay uninitialized: only the refusals, which happen before any launch, are exercised. + """ + x = torch.zeros(tokens, hidden, dtype=torch.bfloat16, device="cuda") + logits = torch.zeros(tokens, NUM_EXPERTS, dtype=torch.float32, device="cuda") + bias = torch.zeros(NUM_EXPERTS, dtype=torch.float32, device="cuda") + w3_w1 = torch.empty(NUM_EXPERTS, 2 * inter, hidden, dtype=torch.bfloat16, device="cuda") + w2 = torch.empty(NUM_EXPERTS, hidden, inter, dtype=torch.bfloat16, device="cuda") + return x, logits, bias, w3_w1, w2 + + +def test_refuses_more_than_one_token() -> None: + """Draft step 0 without the row select has 8 rows: the caller must take another MoE path. + + The op says so rather than computing token 0 alone. + """ + with pytest.raises(RuntimeError, match="takes one token"): + with validating(op): + op(*_small(tokens=2), TOP_K, ROUTED_SCALING_FACTOR) + + +def test_refuses_an_expert_slice_it_does_not_serve() -> None: + """The FC2 kernel walks the expert slice in 64-element K steps.""" + with pytest.raises(RuntimeError, match="unsupported hidden size"): + with validating(op): + op(*_small(inter=96), TOP_K, ROUTED_SCALING_FACTOR) + + +def test_refuses_another_top_k() -> None: + with pytest.raises(RuntimeError, match="top-8"): + with validating(op): + op(*_small(), 6, ROUTED_SCALING_FACTOR) + + +def test_refuses_a_misaligned_token() -> None: + """The FC1 kernel reads x in 16-byte pieces.""" + _, logits, bias, w3_w1, w2 = _small() + storage = torch.zeros(1, 2048 + 1, dtype=torch.bfloat16, device="cuda") + misaligned = storage[:, 1:] + assert misaligned.is_contiguous() and misaligned.data_ptr() % 16 != 0 + with pytest.raises(RuntimeError, match="16-byte aligned"): + with validating(op): + op(misaligned, logits, bias, w3_w1, w2, TOP_K, ROUTED_SCALING_FACTOR) From a32608f5c902eb14083c9490ce348769dd6bf5fb Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:09:47 -0700 Subject: [PATCH 23/53] Add the GLM-5.2 fused verify acceptance op with the forced count torch.ops.trtllm.glm_verify_head_accept computes one verify step's accepted-token counts in one CTA (glmMoeDsa/glmVerifyHeadAccept.cu): a context row accepts its sampled token; a generation row accepts its target token plus the leading draft tokens equal to the targets, and the forced count of TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS then replaces it as the one-model spec worker forms it: the integral part, or that plus one when the pool entry (counter * counter_stride + slot * slot_stride) & (pool_size - 1) is below the fraction (compared in FP32), capped at draft_len + 1, with the counter advanced once per call. With an accepted_tokens buffer it also writes the accepted rows; otherwise the caller reads them from the sampled tokens. int32 and int64 targets, padded (non-ragged) layout. Under PDL it waits for the sampler before reading anything. The tp16 GLM target's owned entry ops/spec/verify_head_accept certifies the c1 verify step (K 7, forced 2.91) and the strict, integral-force, mixed and multi-request forms on sm107 against the ATen chain; its test also drives SpecWorkerBase._apply_force_accepted_tokens on the same pool and counter over consecutive steps, and CUDA-graph replays. Signed-off-by: Dongfeng Yu --- .../kernels/glmMoeDsa/glmVerifyHeadAccept.cu | 122 +++++++++ .../kernels/glmMoeDsa/glmVerifyHeadAccept.h | 72 ++++++ cpp/tensorrt_llm/thop/CMakeLists.txt | 3 +- .../thop/glmVerifyHeadAcceptOp.cpp | 108 ++++++++ .../ops/spec/verify_head_accept.py | 238 ++++++++++++++++++ .../_torch/custom_ops/cpp_custom_ops.py | 10 + .../test_modeling_v2_verify_head_accept.py | 201 +++++++++++++++ 7 files changed, 753 insertions(+), 1 deletion(-) create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.cu create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.h create mode 100644 cpp/tensorrt_llm/thop/glmVerifyHeadAcceptOp.cpp create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/verify_head_accept.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_verify_head_accept.py diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.cu new file mode 100644 index 000000000000..4fd734cab81d --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.cu @@ -0,0 +1,122 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ +namespace +{ + +constexpr int kAcceptBlockSize = 128; + +// One CTA; a thread per request. The forced-acceptance counter is read and advanced once per call. +template +__global__ void __launch_bounds__(kAcceptBlockSize) verifyHeadAcceptKernel(VerifyHeadAcceptParams params) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + // The target tokens are written by the preceding kernel. + cudaGridDependencySynchronize(); +#endif + __shared__ int64_t counter; + bool const randomForce = params.forceMode == VerifyHeadForceMode::kRandom; + if (randomForce && threadIdx.x == 0) + { + counter = *params.forceCounter + 1; + *params.forceCounter = counter; + } + __syncthreads(); + + auto const* target = static_cast(params.targetTokens); + int const batchSize = params.numContexts + params.numGens; + int const rowLen = params.draftLen + 1; + for (int req = static_cast(threadIdx.x); req < batchSize; req += kAcceptBlockSize) + { + int64_t const rowOffset = static_cast(req) * rowLen; + if (req < params.numContexts) + { + params.numAcceptedTokens[req] = 1; + if (params.acceptedTokens != nullptr) + { + params.acceptedTokens[rowOffset] = static_cast(target[req]); + } + continue; + } + int const gen = req - params.numContexts; + TargetT const* genTarget = target + params.numContexts + static_cast(gen) * rowLen; + int32_t const* genDraft = params.draftTokens + static_cast(gen) * params.draftLen; + // The target token plus the leading draft tokens that equal the target tokens. + int accepted = 1; + while (accepted <= params.draftLen + && static_cast(genDraft[accepted - 1]) == static_cast(genTarget[accepted - 1])) + { + ++accepted; + } + if (params.forceMode == VerifyHeadForceMode::kFixed) + { + accepted = params.forceBase; + } + else if (randomForce) + { + int64_t const poolIdx + = (counter * params.counterStride + static_cast(gen) * params.slotStride) & params.poolMask; + int const extra = params.forcePool[poolIdx] < params.forceFrac ? 1 : 0; + accepted = min(params.forceBase + extra, params.forceMax); + } + params.numAcceptedTokens[req] = accepted; + if (params.acceptedTokens != nullptr) + { + for (int j = 0; j < rowLen; ++j) + { + params.acceptedTokens[rowOffset + j] = static_cast(genTarget[j]); + } + } + } +} + +} // namespace + +void invokeVerifyHeadAccept(VerifyHeadAcceptParams const& params) +{ + TLLM_CHECK_WITH_INFO(params.forceMode != VerifyHeadForceMode::kRandom + || (params.forcePool != nullptr && params.forceCounter != nullptr), + "verifyHeadAccept: the random forced count needs the pool and the counter"); + if (params.numContexts + params.numGens == 0) + { + return; + } + if (params.targetTokensInt64) + { + tensorrt_llm::common::launchWithPdlWhenEnabled("verifyHeadAccept", verifyHeadAcceptKernel, dim3(1), + dim3(kAcceptBlockSize), 0, params.stream, params); + } + else + { + tensorrt_llm::common::launchWithPdlWhenEnabled("verifyHeadAccept", verifyHeadAcceptKernel, dim3(1), + dim3(kAcceptBlockSize), 0, params.stream, params); + } +} + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.h new file mode 100644 index 000000000000..52cd9212e7c6 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.h @@ -0,0 +1,72 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +//! How the accepted counts of the generation requests are forced (TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS). +enum class VerifyHeadForceMode : int32_t +{ + kNone = 0, //!< The strict acceptance count + kFixed = 1, //!< forceBase + kRandom = 2, //!< min(forceBase + (pool[(counter * counterStride + gen * slotStride) & poolMask] < forceFrac), + //!< forceMax), counter advanced by one per call +}; + +//! \brief Inputs and outputs of the fused strict + forced acceptance of one-model speculative decoding. +struct VerifyHeadAcceptParams +{ + void const* targetTokens; //!< [numContexts + numGens * (draftLen + 1)], int32, or int64 when targetTokensInt64 + bool targetTokensInt64; + int32_t const* draftTokens; //!< [numGens, draftLen] + int32_t* numAcceptedTokens; //!< Output [numContexts + numGens] + int32_t* acceptedTokens; //!< Output [numContexts + numGens, draftLen + 1], or nullptr: not written + int numContexts; + int numGens; + int draftLen; + VerifyHeadForceMode forceMode; + int forceBase; + int forceMax; + float forceFrac; + float const* forcePool; //!< kRandom: fixed uniform pool + int64_t* forceCounter; //!< kRandom: call counter + int64_t counterStride; + int64_t slotStride; + int64_t poolMask; + cudaStream_t stream; +}; + +//! \brief One kernel for SpecWorkerBase._sample_and_accept_draft_tokens_base after sampling. +//! +//! Context rows accept their sampled token (count 1, accepted row [target, ...]). A generation row accepts its target +//! token plus the leading draft tokens equal to the target tokens (the cumprod count of the strict acceptance), and +//! then the forced count replaces it as forceMode says. Integer results equal the ATen ops bit for bit. Under PDL +//! the kernel waits for its producer before it reads memory and lets its dependents launch at its end. +void invokeVerifyHeadAccept(VerifyHeadAcceptParams const& params); + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index d50b691a6b46..8a84a5bf4724 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -165,7 +165,8 @@ add_th_common_family( loraOp.cpp inplaceSliceCopyOp.cpp mhcOp.cpp - compressorOp.cpp) + compressorOp.cpp + glmVerifyHeadAcceptOp.cpp) add_library( th_common SHARED diff --git a/cpp/tensorrt_llm/thop/glmVerifyHeadAcceptOp.cpp b/cpp/tensorrt_llm/thop/glmVerifyHeadAcceptOp.cpp new file mode 100644 index 000000000000..cfeb6d2ea623 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmVerifyHeadAcceptOp.cpp @@ -0,0 +1,108 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +// Accepted-token counts of the strict acceptance of one-model speculative decoding, with the forced count of +// TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS applied (force_mode 0: none, 1: force_base, 2: force_base plus a pool +// draw below force_frac, capped at force_max, force_counter advanced by one). Writes the accepted-token rows into +// accepted_tokens when given. +at::Tensor glmVerifyHeadAccept(at::Tensor const& target_tokens, at::Tensor const& draft_tokens, int64_t num_contexts, + std::optional accepted_tokens, std::optional const& force_pool, + std::optional force_counter, int64_t force_mode, int64_t force_base, int64_t force_max, + double force_frac, int64_t counter_stride, int64_t slot_stride, int64_t pool_size) +{ + CHECK_TH_CUDA(target_tokens); + CHECK_CONTIGUOUS(target_tokens); + auto const targetType = target_tokens.scalar_type(); + TORCH_CHECK(targetType == at::ScalarType::Int || targetType == at::ScalarType::Long, + "glm_verify_head_accept expects int32 or int64 target tokens, got ", targetType); + CHECK_INPUT(draft_tokens, at::ScalarType::Int); + TORCH_CHECK(draft_tokens.dim() == 2, "glm_verify_head_accept expects draft tokens [num_gens, draft_len]"); + int64_t const numGens = draft_tokens.size(0); + int64_t const draftLen = draft_tokens.size(1); + int64_t const batchSize = num_contexts + numGens; + TORCH_CHECK(target_tokens.numel() >= num_contexts + numGens * (draftLen + 1), + "glm_verify_head_accept: target tokens hold one row per context and draft_len + 1 per generation request"); + CHECK_OPTIONAL_INPUT(accepted_tokens, at::ScalarType::Int); + TORCH_CHECK(!accepted_tokens.has_value() || accepted_tokens->numel() >= batchSize * (draftLen + 1), + "glm_verify_head_accept: accepted_tokens must be [batch, draft_len + 1]"); + TORCH_CHECK(force_mode >= 0 && force_mode <= 2, "glm_verify_head_accept: unknown force_mode ", force_mode); + bool const randomForce = force_mode == 2; + if (randomForce) + { + TORCH_CHECK(force_pool.has_value() && force_counter.has_value(), + "glm_verify_head_accept: force_mode 2 needs force_pool and force_counter"); + CHECK_INPUT(force_pool.value(), at::ScalarType::Float); + CHECK_INPUT(force_counter.value(), at::ScalarType::Long); + TORCH_CHECK(pool_size > 0 && (pool_size & (pool_size - 1)) == 0 && force_pool->numel() >= pool_size, + "glm_verify_head_accept: the pool size must be a power of two within force_pool"); + } + + at::Tensor numAccepted = at::empty({batchSize}, draft_tokens.options()); + + namespace tk = tensorrt_llm::kernels::glm_moe_dsa; + tk::VerifyHeadAcceptParams params; + params.targetTokens = target_tokens.data_ptr(); + params.targetTokensInt64 = targetType == at::ScalarType::Long; + params.draftTokens = draft_tokens.data_ptr(); + params.numAcceptedTokens = numAccepted.data_ptr(); + params.acceptedTokens = accepted_tokens.has_value() ? accepted_tokens->data_ptr() : nullptr; + params.numContexts = static_cast(num_contexts); + params.numGens = static_cast(numGens); + params.draftLen = static_cast(draftLen); + params.forceMode = static_cast(force_mode); + params.forceBase = static_cast(force_base); + params.forceMax = static_cast(force_max); + params.forceFrac = static_cast(force_frac); + params.forcePool = randomForce ? force_pool->data_ptr() : nullptr; + params.forceCounter = randomForce ? force_counter->data_ptr() : nullptr; + params.counterStride = counter_stride; + params.slotStride = slot_stride; + params.poolMask = pool_size - 1; + params.stream = at::cuda::getCurrentCUDAStream(draft_tokens.get_device()); + tk::invokeVerifyHeadAccept(params); + return numAccepted; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_verify_head_accept(Tensor target_tokens, Tensor draft_tokens, int num_contexts, " + "Tensor(a!)? accepted_tokens, Tensor? force_pool, Tensor(b!)? force_counter, int force_mode, int force_base, " + "int force_max, float force_frac, int counter_stride, int slot_stride, int pool_size) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_verify_head_accept", &tensorrt_llm::torch_ext::glmVerifyHeadAccept); +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/verify_head_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/verify_head_accept.py new file mode 100644 index 000000000000..e8acff1335e4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/verify_head_accept.py @@ -0,0 +1,238 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Strict draft acceptance with the forced accepted-token count in one kernel, via the glm_verify_head_accept op.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class GlmVerifyHeadAccept(OpWrapper): + """The accepted-token counts of one verify step, and optionally its accepted-token rows. + + A context request accepts its one sampled token (count 1). A generation request with sampled target tokens + `t[0..K]` and draft tokens `d[0..K-1]` accepts `1 + (length of the leading run of d[j] == t[j])`, then the + forced count replaces that as `force_mode` says: + - 0: no force; + - 1: `force_base`; + - 2: `min(force_base + (force_pool[i] < force_frac), force_max)` with + `i = (c * counter_stride + g * slot_stride) & (pool_size - 1)`, `g` the generation index and `c` the value + of `force_counter` after this call adds one to it (the comparison in FP32). + + That is `SpecWorkerBase._sample_and_accept_draft_tokens_base` after sampling with + `_apply_force_accepted_tokens`, for the padded layout (`target_tokens` holds one token per context request, + then `K + 1` per generation request). Returns the counts `[num_contexts + num_gens]` int32. With + `accepted_tokens` the op also writes the accepted-token rows: column 0 of each context row, the whole row of + each generation request (its target tokens); without it the caller reads the generation rows from + `target_tokens` itself. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # spec: num_contexts, num_gens, draft_len (K = 7 at c1), target dtype, matches (the leading draft tokens equal + # to the targets, per generation request), force (mode, base, max, frac), rows (accepted_tokens given). The + # hash constants are the spec worker's: counter stride 6007, slot stride 1009, pool 2^16. + # At c1 the gen server verifies one request with K = 7 and the cell config forces 2.91 accepted draft tokens: + # mode 2, base 3, max 8, frac 0.91; the accepted rows alias the sampled tokens. + CELLS: tuple[Cell, ...] = ( + Cell( + why="c1 verify step: one request, K 7, forced 2.91 (base 3, frac 0.91)", + spec=dict( + num_contexts=0, + num_gens=1, + draft_len=7, + target_dtype=torch.int32, + matches=(2,), + force=(2, 3, 8, 0.91), + rows=False, + ), + ), + Cell( + why="c1 verify step without a forced count: the strict count", + spec=dict( + num_contexts=0, + num_gens=1, + draft_len=7, + target_dtype=torch.int32, + matches=(5,), + force=(0, 0, 0, 0.0), + rows=False, + ), + ), + Cell( + why="every draft token accepted: the count reaches K + 1", + spec=dict( + num_contexts=0, + num_gens=1, + draft_len=7, + target_dtype=torch.int32, + matches=(7,), + force=(0, 0, 0, 0.0), + rows=True, + ), + ), + Cell( + why="an integral forced count (3.0): mode 1, every request gets base 4", + spec=dict( + num_contexts=0, + num_gens=1, + draft_len=7, + target_dtype=torch.int32, + matches=(0,), + force=(1, 4, 8, 0.0), + rows=False, + ), + ), + Cell( + why="warm-up batch with a context request, int64 targets, rows written", + spec=dict( + num_contexts=1, + num_gens=1, + draft_len=7, + target_dtype=torch.int64, + matches=(3,), + force=(2, 3, 8, 0.91), + rows=True, + ), + ), + Cell( + why="several generation requests: the pool index follows the slot", + spec=dict( + num_contexts=0, + num_gens=4, + draft_len=7, + target_dtype=torch.int32, + matches=(0, 7, 2, 4), + force=(2, 3, 8, 0.5), + rows=True, + ), + ), + ) + + note = """ + Not for ragged verification: a request that verified fewer than K + 1 positions has no padded row, and the op + would read its neighbour's tokens. The caller keeps the ATen chain there. + + The kernel waits for its producer (the sampler) before it reads anything and lets its dependents launch at its + end. Mode 2 advances `force_counter` once per call, also inside a captured graph's replays, exactly as the ATen + chain's in-place add does, so the op and the chain draw the same pool entries given the same counter. + """ + + def __call__( + self, + target_tokens: torch.Tensor, + draft_tokens: torch.Tensor, + num_contexts: int, + accepted_tokens: Optional[torch.Tensor], + force_pool: Optional[torch.Tensor], + force_counter: Optional[torch.Tensor], + force_mode: int, + force_base: int, + force_max: int, + force_frac: float, + counter_stride: int, + slot_stride: int, + pool_size: int, + ) -> torch.Tensor: + return torch.ops.trtllm.glm_verify_head_accept( + target_tokens, + draft_tokens, + num_contexts, + accepted_tokens, + force_pool, + force_counter, + force_mode, + force_base, + force_max, + force_frac, + counter_stride, + slot_stride, + pool_size, + ) + + def reference( + self, + target_tokens: torch.Tensor, + draft_tokens: torch.Tensor, + num_contexts: int, + accepted_tokens: Optional[torch.Tensor], + force_pool: Optional[torch.Tensor], + force_counter: Optional[torch.Tensor], + force_mode: int, + force_base: int, + force_max: int, + force_frac: float, + counter_stride: int, + slot_stride: int, + pool_size: int, + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + """The ATen chain of the one-model spec worker: `(counts, rows, counter)`. + + `rows` is what `accepted_tokens` holds after the call (its unwritten entries keep the values it was handed) + and `counter` what `force_counter` holds after it; each is None when the op is not handed that buffer + (or, for the counter, does not advance it). Both buffers are read as they were before the call. + """ + num_gens, draft_len = draft_tokens.shape + batch_size = num_contexts + num_gens + target = target_tokens.reshape(-1) + counts = torch.ones(batch_size, dtype=torch.int32, device=draft_tokens.device) + gen_target = target[num_contexts : num_contexts + num_gens * (draft_len + 1)].reshape( + num_gens, draft_len + 1 + ) + counts[num_contexts:] += ( + torch.cumprod((draft_tokens.long() == gen_target[:, :draft_len].long()).int(), dim=-1) + .sum(1) + .int() + ) + counter = None + if force_mode == 1: + counts[num_contexts:] = force_base + elif force_mode == 2: + counter = force_counter + 1 + slots = torch.arange(num_gens, dtype=torch.int64, device=draft_tokens.device) + index = (counter * counter_stride + slots * slot_stride) & (pool_size - 1) + extra = (force_pool[index] < force_frac).int() + counts[num_contexts:] = (force_base + extra).clamp(max=force_max).int() + rows = None + if accepted_tokens is not None: + rows = accepted_tokens.clone() + rows[:num_contexts, 0] = target[:num_contexts].int() + rows[num_contexts:batch_size] = gen_target.int() + return counts, rows, counter + + def is_valid( + self, + target_tokens: torch.Tensor, + draft_tokens: torch.Tensor, + num_contexts: int, + accepted_tokens: Optional[torch.Tensor], + force_pool: Optional[torch.Tensor], + force_counter: Optional[torch.Tensor], + force_mode: int, + force_base: int, + force_max: int, + force_frac: float, + counter_stride: int, + slot_stride: int, + pool_size: int, + ) -> None: + num_gens, draft_len = draft_tokens.shape + # The kernel writes row r at offset r * (K + 1): a wider buffer (rows of the maximum draft length while + # this step verifies fewer) is accepted by the op's size check and filled with the wrong row stride. + assert accepted_tokens is None or tuple(accepted_tokens.shape) == ( + num_contexts + num_gens, + draft_len + 1, + ), "accepted_tokens must be exactly [num_contexts + num_gens, draft_len + 1]" + # A forced count outside [1, K + 1] is written as given and sends the caller past its draft window. + if force_mode != 0: + assert 1 <= force_base <= force_max == draft_len + 1, ( + "a forced count needs 1 <= force_base <= force_max == draft_len + 1" + ) + + +glm_verify_head_accept = GlmVerifyHeadAccept() diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index fea831399cb6..bf24256133dd 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -939,6 +939,16 @@ def _( logits.new_empty(logits.shape, dtype=torch.float32), ) + @torch.library.register_fake("trtllm::glm_verify_head_accept") + def _(target_tokens: torch.Tensor, draft_tokens: torch.Tensor, + num_contexts: int, accepted_tokens: Optional[torch.Tensor], + force_pool: Optional[torch.Tensor], + force_counter: Optional[torch.Tensor], force_mode: int, + force_base: int, force_max: int, force_frac: float, + counter_stride: int, slot_stride: int, pool_size: int): + return draft_tokens.new_empty( + (num_contexts + draft_tokens.shape[0], )) + @torch.library.register_fake("trtllm::fused_compute_probs_from_logits") def _( logits: torch.Tensor, diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_verify_head_accept.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_verify_head_accept.py new file mode 100644 index 000000000000..6aa319935f24 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_verify_head_accept.py @@ -0,0 +1,201 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm_verify_head_accept entry.""" + +import functools +import zlib +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec.verify_head_accept import ( # noqa: E501 + glm_verify_head_accept as op, +) +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + SpecWorkerBase, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" +) + +VOCAB = 154880 +STRIDES = ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, +) + + +def _spec_worker_rng_state(device): + """The pool and counter the one-model spec worker builds (fixed seed, identical on every rank).""" + holder = SimpleNamespace(_force_accept_rng_pool=None, _force_accept_rng_counter=None) + SpecWorkerBase._ensure_force_accept_rng_state(holder, device) + return holder._force_accept_rng_pool, holder._force_accept_rng_counter + + +def _build(spec, seed): + """One cell's operands. + + Target tokens, draft tokens with `matches[g]` leading tokens equal to request g's targets, the spec worker's + pool, a counter at a seed-dependent value and, with `rows`, a sentinel-filled accepted_tokens. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + num_contexts, num_gens, draft_len = spec["num_contexts"], spec["num_gens"], spec["draft_len"] + num_targets = num_contexts + num_gens * (draft_len + 1) + target = torch.randint(0, VOCAB - 1, (num_targets,), generator=g, device="cuda") + gen_target = target[num_contexts:].view(num_gens, draft_len + 1) + draft = torch.randint(0, VOCAB, (num_gens, draft_len), generator=g, device="cuda") + for request, matched in enumerate(spec["matches"]): + draft[request, :matched] = gen_target[request, :matched] + if matched < draft_len: + # Targets stay below VOCAB - 1, so this is a valid token that differs from the target. + draft[request, matched] = gen_target[request, matched] + 1 + pool, counter = _spec_worker_rng_state(torch.device("cuda")) + counter.fill_(seed % 997) + rows = None + if spec["rows"]: + rows = torch.full( + (num_contexts + num_gens, draft_len + 1), -7, dtype=torch.int32, device="cuda" + ) + return ( + target.to(spec["target_dtype"]).contiguous(), + draft.to(torch.int32).contiguous(), + rows, + pool, + counter, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: + """Integer outputs: the op and the ATen chain agree exactly.""" + spec = cell.spec + target, draft, rows, pool, counter = _build(spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + mode, base, maximum, frac = spec["force"] + rows_before = rows.clone() if rows is not None else None + counter_before = counter.clone() + certifying(op) + contexts = spec["num_contexts"] + counts = op(target, draft, contexts, rows, pool, counter, mode, base, maximum, frac, *STRIDES) + ref_counts, ref_rows, ref_counter = op.reference( + target, + draft, + spec["num_contexts"], + rows_before, + pool, + counter_before, + mode, + base, + maximum, + frac, + *STRIDES, + ) + torch.testing.assert_close(counts, ref_counts, rtol=0, atol=0) + if rows is not None: + torch.testing.assert_close(rows, ref_rows, rtol=0, atol=0) + if ref_counter is not None: + torch.testing.assert_close(counter, ref_counter, rtol=0, atol=0) + else: + torch.testing.assert_close(counter, counter_before, rtol=0, atol=0) + + +@pytest.mark.parametrize("forced", [2.91, 2.5, 6.1, 3.0]) +def test_forced_count_matches_the_spec_worker(forced) -> None: + """The op's forced counts equal the spec worker's over consecutive steps. + + `SpecWorkerBase._apply_force_accepted_tokens` runs on its own counter, the op on a copy of it with the same + pool and the arguments the target derives from the forced value as that method does. + """ + draft_len, num_gens = 7, 4 + worker = SimpleNamespace( + force_num_accepted_tokens=forced, + _force_accept_rng_pool=None, + _force_accept_rng_counter=None, + ) + worker._ensure_force_accept_rng_state = functools.partial( + SpecWorkerBase._ensure_force_accept_rng_state, worker + ) + worker._ensure_force_accept_rng_state(torch.device("cuda")) + op_counter = worker._force_accept_rng_counter.clone() + int_part = int(forced) + frac_part = forced - int_part + max_total = draft_len + 1 + base_total = min(int_part + 1, max_total) + mode = 2 if frac_part > 0.0 and base_total < max_total else 1 + target, draft, _, _, _ = _build( + dict( + num_contexts=0, + num_gens=num_gens, + draft_len=draft_len, + target_dtype=torch.int32, + matches=(0,) * num_gens, + rows=False, + ), + seed=17, + ) + for _ in range(64): + strict = torch.ones(num_gens, dtype=torch.int32, device="cuda") + expected = SpecWorkerBase._apply_force_accepted_tokens(worker, strict, 0, draft_len) + counts = op( + target, + draft, + 0, + None, + worker._force_accept_rng_pool, + op_counter, + mode, + base_total, + max_total, + frac_part if mode == 2 else 0.0, + *STRIDES, + ) + torch.testing.assert_close(counts, expected, rtol=0, atol=0) + torch.testing.assert_close(op_counter, worker._force_accept_rng_counter, rtol=0, atol=0) + + +def test_cuda_graph_replays_advance_the_counter() -> None: + """Inside the decode graph every replay draws the next pool entries, as the eager calls would.""" + spec = op.CELLS[0].spec + target, draft, _, pool, counter = _build(spec, seed=5) + force = spec["force"] + # One eager call first, as the engine's warm-up does before capture. + op(target, draft, 0, None, pool, counter, *force, *STRIDES) + expected_counter = counter.clone() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_counts = op(target, draft, 0, None, pool, counter, *force, *STRIDES) + for _ in range(8): + expected, _, expected_counter = op.reference( + target, draft, 0, None, pool, expected_counter, *force, *STRIDES + ) + graph.replay() + torch.testing.assert_close(graph_counts, expected, rtol=0, atol=0) + torch.testing.assert_close(counter, expected_counter, rtol=0, atol=0) + + +def test_guard_refuses_a_wider_rows_buffer() -> None: + """Rows of the maximum draft length on a step that verifies fewer. + + The op's size check accepts them, and the rows are written with the wrong row stride. + """ + target, draft, _, pool, counter = _build(op.CELLS[0].spec, seed=3) + wide = torch.zeros(1, 12, dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match="exactly"): + with validating(op): + op(target, draft, 0, wide, pool, counter, 0, 0, 0, 0.0, *STRIDES) + + +def test_guard_refuses_a_forced_count_past_the_draft_window() -> None: + target, draft, _, pool, counter = _build(op.CELLS[0].spec, seed=4) + with pytest.raises(AssertionError, match="force_max == draft_len"): + with validating(op): + op(target, draft, 0, None, pool, counter, 1, 9, 9, 0.0, *STRIDES) From 80634f6f8e4aaff98487b7435b05cdaa15f33cec Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:09:57 -0700 Subject: [PATCH 24/53] Add the GLM-5.2 MTP step-0 row selection as one Triton launch ops/spec/gather_last_accepted_rows selects, for each request r, row g[r] = batch_indices[r] * rows_per_request + num_accepted[r] - 1 + row_offset (its last accepted token) of one or two row tensors and returns g as int64 with the gathered rows (_gather_last_accepted_rows._gather_last_accepted_rows_kernel). Draft step 0 of the tp16 GLM target's MTP loop uses it twice over: to take the target's hidden rows and gather_ids in one launch, and, after the MTP layer's post-attention all-reduce, to keep only the selected rows of the normalized hidden states and the residual so that the MoE, its all-reduce and the draft LM head run on one row per request. The entry certifies the c1 shapes (8 verified rows of 6144, with and without the residual) and a permuted 4-request batch bit for bit on sm107; its guard refuses counts outside [1, rows_per_request], strided index vectors and a second tensor of another width. Signed-off-by: Dongfeng Yu --- .../ops/spec/_gather_last_accepted_rows.py | 99 ++++++++++++++ .../ops/spec/gather_last_accepted_rows.py | 123 ++++++++++++++++++ ...t_modeling_v2_gather_last_accepted_rows.py | 88 +++++++++++++ 3 files changed, 310 insertions(+) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_gather_last_accepted_rows.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/gather_last_accepted_rows.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_gather_last_accepted_rows.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_gather_last_accepted_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_gather_last_accepted_rows.py new file mode 100644 index 000000000000..492124cbbf1d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_gather_last_accepted_rows.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Triton kernel behind ``gather_last_accepted_rows``: the step-0 row selection of the MTP draft loop. + +Request ``r`` selects row ``g[r] = batch_indices[r] * rows_per_request + num_accepted[r] - 1 + row_offset``, the +row of its last accepted token. One launch writes the int64 ``g`` and copies row ``g[r]`` of one or two row +tensors into row ``r`` of fresh outputs. +""" + +from __future__ import annotations + +from typing import Optional, Tuple, Union + +import torch +import triton +import triton.language as tl + +#: Columns one program copies. +_BLOCK = 1024 + + +@triton.jit +def _gather_last_accepted_rows_kernel( + src_ptr, + out_ptr, + src2_ptr, + out2_ptr, + batch_indices_ptr, + num_accepted_ptr, + gather_ids_ptr, + stride_src, + stride_out, + stride_src2, + stride_out2, + rows_per_request, + row_offset, + n_cols, + HAS_SECOND: tl.constexpr, + BLOCK: tl.constexpr, +): + req = tl.program_id(0) + chunk = tl.program_id(1) + row = ( + tl.load(batch_indices_ptr + req).to(tl.int64) * rows_per_request + + tl.load(num_accepted_ptr + req).to(tl.int64) + - 1 + + row_offset + ) + if chunk == 0: + tl.store(gather_ids_ptr + req, row) + cols = chunk * BLOCK + tl.arange(0, BLOCK) + mask = cols < n_cols + values = tl.load(src_ptr + row * stride_src + cols, mask=mask) + tl.store(out_ptr + req.to(tl.int64) * stride_out + cols, values, mask=mask) + if HAS_SECOND: + values2 = tl.load(src2_ptr + row * stride_src2 + cols, mask=mask) + tl.store(out2_ptr + req.to(tl.int64) * stride_out2 + cols, values2, mask=mask) + + +def gather_last_accepted_rows( + rows: torch.Tensor, + batch_indices: torch.Tensor, + num_accepted: torch.Tensor, + rows_per_request: int, + row_offset: int, + second: Optional[torch.Tensor] = None, +) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + """``(g, rows[g])``, plus ``second[g]`` as a third element when ``second`` is given, in one launch. + + ``rows`` and ``second`` are 2D with unit column stride; ``batch_indices`` and ``num_accepted`` are contiguous + integer vectors of one length. + """ + num_requests = batch_indices.shape[0] + n_cols = rows.shape[1] + gather_ids = torch.empty((num_requests,), dtype=torch.int64, device=rows.device) + out = rows.new_empty((num_requests, n_cols)) + out2 = second.new_empty((num_requests, second.shape[1])) if second is not None else out + if num_requests > 0: + _gather_last_accepted_rows_kernel[(num_requests, triton.cdiv(n_cols, _BLOCK))]( + rows, + out, + second if second is not None else rows, + out2, + batch_indices, + num_accepted, + gather_ids, + rows.stride(0), + out.stride(0), + second.stride(0) if second is not None else 0, + out2.stride(0), + rows_per_request, + row_offset, + n_cols, + HAS_SECOND=second is not None, + BLOCK=_BLOCK, + ) + if second is not None: + return gather_ids, out, out2 + return gather_ids, out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/gather_last_accepted_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/gather_last_accepted_rows.py new file mode 100644 index 000000000000..d84da2985147 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/gather_last_accepted_rows.py @@ -0,0 +1,123 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MTP draft step 0's row selection -- each request's last accepted row, and its index -- in one Triton launch.""" + +from __future__ import annotations + +from typing import Optional, Tuple, Union + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._gather_last_accepted_rows import gather_last_accepted_rows as _gather_last_accepted_rows + + +class GlmGatherLastAcceptedRows(OpWrapper): + """Each request's last accepted row of `rows` (and of `second`), and the row's index. + + Request `r` selects `g[r] = batch_indices[r] * rows_per_request + num_accepted[r] - 1 + row_offset`. Draft + step 0 of the MTP loop runs the layer on every verified row of a request (`K + 1` of them), but only the row of + the request's last accepted token feeds the LM head and the next step. Returns `(g, rows[g])`, or + `(g, rows[g], second[g])`: `g` int64 `[R]`, the rows fresh `[R, n]` tensors. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # At c1 one request verifies K + 1 = 8 rows and step 0 selects one: after the post-attention all-reduce, the + # normalized hidden rows and the residual together (the MoE, its all-reduce and the LM head then run on that + # row), or the target's hidden rows alone where no row select applies. Generation-only batches only, so the + # row offset (the context tokens ahead of them) is 0. + CELLS: tuple[Cell, ...] = ( + Cell( + why="c1 step-0 row select: hidden and residual rows [8, 6144], one request", + spec=dict(requests=1, rows_per_request=8, cols=6144, second=True, accepted=(3,)), + ), + Cell( + why="c1 step-0 gather of the target hidden rows alone, every draft accepted", + spec=dict(requests=1, rows_per_request=8, cols=6144, second=False, accepted=(8,)), + ), + Cell( + why="four requests in a permuted batch order, first rows accepted", + spec=dict( + requests=4, + rows_per_request=8, + cols=6144, + second=True, + accepted=(1, 8, 4, 2), + batch_order=(2, 0, 3, 1), + ), + ), + ) + + note = """ + The kernel reads `batch_indices` and `num_accepted` as dense vectors and does not bound `g`: a count outside + [1, rows_per_request] reads another request's row or memory outside `rows`. + """ + + def __call__( + self, + rows: torch.Tensor, + batch_indices: torch.Tensor, + num_accepted: torch.Tensor, + rows_per_request: int, + row_offset: int, + second: Optional[torch.Tensor] = None, + ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + return _gather_last_accepted_rows( + rows, batch_indices, num_accepted, rows_per_request, row_offset, second + ) + + def reference( + self, + rows: torch.Tensor, + batch_indices: torch.Tensor, + num_accepted: torch.Tensor, + rows_per_request: int, + row_offset: int, + second: Optional[torch.Tensor] = None, + ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + """The draft loop's index chain and its row gathers.""" + gather_ids = batch_indices.long() * rows_per_request + num_accepted.long() - 1 + row_offset + if second is None: + return gather_ids, rows[gather_ids] + return gather_ids, rows[gather_ids], second[gather_ids] + + def is_valid( + self, + rows: torch.Tensor, + batch_indices: torch.Tensor, + num_accepted: torch.Tensor, + rows_per_request: int, + row_offset: int, + second: Optional[torch.Tensor] = None, + ) -> None: + # The kernel addresses rows through stride(0) with unit column stride, and the index vectors as dense + # arrays; any other layout is read from the wrong elements. + assert rows.dim() == 2 and rows.stride(1) == 1, "rows must be 2D with unit column stride" + assert second is None or (second.dim() == 2 and second.stride(1) == 1), ( + "second must be 2D with unit column stride" + ) + # The kernel copies rows.shape[1] columns of both tensors: a wider second keeps garbage in its tail. + assert second is None or second.shape[1] == rows.shape[1], ( + "second must have the columns of rows" + ) + assert batch_indices.dim() == 1 and batch_indices.is_contiguous(), ( + "batch_indices must be a contiguous vector" + ) + assert num_accepted.dim() == 1 and num_accepted.is_contiguous(), ( + "num_accepted must be a contiguous vector" + ) + assert num_accepted.shape[0] >= batch_indices.shape[0], "one accepted count per request" + # The selected row is not bounded on the device. + count = num_accepted[: batch_indices.shape[0]] + assert bool(((count >= 1) & (count <= rows_per_request)).all()), ( + "num_accepted must lie in [1, rows_per_request]" + ) + gather_ids = batch_indices.long() * rows_per_request + count.long() - 1 + row_offset + assert bool(((gather_ids >= 0) & (gather_ids < rows.shape[0])).all()), ( + "the selected rows must lie inside rows" + ) + + +glm_gather_last_accepted_rows = GlmGatherLastAcceptedRows() diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_gather_last_accepted_rows.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_gather_last_accepted_rows.py new file mode 100644 index 000000000000..c1e5844278c0 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_gather_last_accepted_rows.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm_gather_last_accepted_rows entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec.gather_last_accepted_rows import ( # noqa: E501 + glm_gather_last_accepted_rows as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" +) + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + requests, per_request, cols = spec["requests"], spec["rows_per_request"], spec["cols"] + rows = torch.randn(requests * per_request, cols, generator=g, device="cuda").to(torch.bfloat16) + second = None + if spec["second"]: + second = torch.randn(requests * per_request, cols, generator=g, device="cuda").to( + torch.bfloat16 + ) + order = spec.get("batch_order", tuple(range(requests))) + batch_indices = torch.tensor(order, dtype=torch.int32, device="cuda") + num_accepted = torch.tensor(spec["accepted"], dtype=torch.int32, device="cuda") + return rows, batch_indices, num_accepted, per_request, 0, second + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: + """A gather moves bytes: the op and the reference agree bit for bit.""" + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + out = op(*args) + ref = op.reference(*args) + assert len(out) == len(ref) == (3 if cell.spec["second"] else 2) + for got, expected in zip(out, ref): + torch.testing.assert_close(got, expected, rtol=0, atol=0) + + +def test_row_offset_shifts_the_selection() -> None: + """Rows after `row_offset` context tokens: the index moves by the offset, the gathers with it.""" + rows, batch_indices, num_accepted, per_request, _, second = _build(op.CELLS[0].spec, seed=3) + padded = torch.cat([torch.zeros(5, rows.shape[1], dtype=rows.dtype, device="cuda"), rows]) + padded_second = torch.cat([torch.zeros_like(padded[:5]), second]) + with validating(op): + out = op(padded, batch_indices, num_accepted, per_request, 5, padded_second) + ref = op.reference(rows, batch_indices, num_accepted, per_request, 0, second) + torch.testing.assert_close(out[0], ref[0] + 5, rtol=0, atol=0) + torch.testing.assert_close(out[1], ref[1], rtol=0, atol=0) + torch.testing.assert_close(out[2], ref[2], rtol=0, atol=0) + + +def test_guard_refuses_a_count_of_zero() -> None: + """A request with no accepted row would select the row before its own.""" + rows, batch_indices, _, per_request, offset, second = _build(op.CELLS[0].spec, seed=4) + zero = torch.zeros(1, dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match=r"\[1, rows_per_request\]"): + with validating(op): + op(rows, batch_indices, zero, per_request, offset, second) + + +def test_guard_refuses_a_wider_second() -> None: + """The kernel copies the columns of `rows`: a wider `second` keeps garbage in its tail.""" + rows, batch_indices, num_accepted, per_request, offset, _ = _build(op.CELLS[0].spec, seed=5) + wide = torch.zeros(rows.shape[0], rows.shape[1] + 8, dtype=rows.dtype, device="cuda") + with pytest.raises(AssertionError, match="columns of rows"): + with validating(op): + op(rows, batch_indices, num_accepted, per_request, offset, wide) + + +def test_guard_refuses_a_strided_count_vector() -> None: + """The kernel reads the counts as a dense vector.""" + rows, _, _, per_request, offset, second = _build(op.CELLS[2].spec, seed=6) + batch_indices = torch.arange(4, dtype=torch.int32, device="cuda") + strided = torch.ones(8, dtype=torch.int32, device="cuda")[::2] + with pytest.raises(AssertionError, match="contiguous vector"): + with validating(op): + op(rows, batch_indices, strided, per_request, offset, second) From 35b6fe11dcb4acb8244e3be592443a26142ccc50 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:10:07 -0700 Subject: [PATCH 25/53] Add the GLM-5.2 MTP draft-step input glue as one CuTe DSL kernel ops/spec/draft_glue_rmsnorm_cat returns this rank's window [col_lo, col_hi) of cat([enorm(embed_table[input_ids]), hnorm(hidden_states)], -1) as one contiguous BF16 tensor, from one kernel (_draft_glue_rmsnorm.DraftGlueRMSNormKernel): it gathers the embedding rows itself (a zero row for an id outside the table), normalizes only the halves the window overlaps and stores only the window. Each row is normalized with FlashInfer's CuTe RMSNormKernel thread geometry, reduction and arithmetic. Under PDL it waits before any load and releases its dependent (eh_proj) right after the wait. At tp16 the eh_proj input window of rank r is 768 columns, so ranks 0-7 normalize only the embedding half and ranks 8-15 only the hidden half; the contiguous window lets eh_proj run on the low-M GEMM at draft step 0's eight rows. The entry certifies the c1 windows of ranks 0, 7, 8 and 15, out-of-table ids and a single row on sm107 against an FP32 reference (and, where FlashInfer runs its CuTe RMSNorm, against flashinfer.norm. rmsnorm bit for bit); the guard refuses windows that are not 8-column aligned and misaligned rows. Signed-off-by: Dongfeng Yu --- .../ops/spec/_draft_glue_rmsnorm.py | 457 ++++++++++++++++++ .../ops/spec/draft_glue_rmsnorm_cat.py | 165 +++++++ ...test_modeling_v2_draft_glue_rmsnorm_cat.py | 101 ++++ 3 files changed, 723 insertions(+) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_glue_rmsnorm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_glue_rmsnorm_cat.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_glue_rmsnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_glue_rmsnorm.py new file mode 100644 index 000000000000..6f4a3ba08625 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_glue_rmsnorm.py @@ -0,0 +1,457 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CuTe DSL kernel behind ``draft_glue_rmsnorm_cat``: the MTP draft-step eh_proj input in one kernel. + +The MTP layer feeds ``eh_proj`` with this rank's columns of ``cat([enorm(embed(input_ids)), hnorm(hidden_states)], +-1)``. ``DraftGlueRMSNormKernel`` gathers the embedding rows from the table itself (a zero row for an id outside the +table), normalizes them with the enorm weight and the hidden rows with the hnorm weight, and stores only the +requested window of the concatenated rows, contiguously. + +Each row is normalized as FlashInfer's CuTe ``RMSNormKernel`` (the kernel ``flashinfer.norm.rmsnorm`` runs for 2D +BF16 inputs) does it: the thread geometry is taken from an ``RMSNormKernel`` instance (64 threads per row for +H = 6144, 8-element vectors, the same column order per thread), the reduction is FlashInfer's +``row_reduce_sum_multirow``, the arithmetic statements are FlashInfer's, and the kernel is compiled with the same +options. The per-row code follows FlashInfer's ``QKRMSNormKernel``, which gives every row its own tile so that a +row can come from any tensor. +""" + +from __future__ import annotations + +import functools + +import cutlass +import cutlass.cute as cute +import torch +from cutlass import Float32, Int64 +from flashinfer.norm.kernels.rmsnorm import RMSNormKernel +from flashinfer.norm.utils import make_tv_layout, predicate_k, row_reduce_sum_multirow + +_IDS_DTYPES = {torch.int32: cutlass.Int32, torch.int64: cutlass.Int64} + + +@cute.jit +def _predicate_cols( + tXcX: cute.Tensor, lo: cutlass.Constexpr[int], hi: cutlass.Constexpr[int] +) -> cute.Tensor: + """Copy predicate that is true where a vector's first column lies in [lo, hi).""" + tXpX = cute.make_rmem_tensor( + cute.make_layout( + ( + cute.size(tXcX, mode=[0, 1]), + cute.size(tXcX, mode=[1]), + cute.size(tXcX, mode=[2]), + ), + stride=(cute.size(tXcX, mode=[2]), 0, 1), + ), + cutlass.Boolean, + ) + for rest_v in cutlass.range_constexpr(tXpX.shape[0]): + for rest_k in cutlass.range_constexpr(tXpX.shape[2]): + col = tXcX[(0, rest_v), 0, rest_k][1] + in_window = cute.elem_less(col, hi) + if cutlass.const_expr(lo > 0): + in_window = in_window & cute.elem_less(lo - 1, col) + tXpX[rest_v, 0, rest_k] = in_window + return tXpX + + +class DraftGlueRMSNormKernel: + """RMSNorm of the embedding rows (half 0) and of the hidden rows (half 1) of the MTP draft step. + + Column ``k`` of half ``h`` is column ``h * H + k`` of the concatenated row. The kernel computes the halves + overlapping the column window ``[col_lo, col_hi)`` and stores the window's columns at + ``out[row, col - col_lo]``. Half 0 reads ``table[input_ids[row]]`` (a zero row for an id outside the table). Each + group of ``threads_per_row`` threads normalizes one (row, half) pair; with both halves a CTA covers one row. + + Args: + hidden_size: H, the length of each normalized row. + half_begin: First half overlapping the window. + num_halves: Number of halves overlapping the window (1 or 2). + col_lo: First concatenated column stored. + col_hi: One past the last concatenated column stored. + """ + + def __init__( + self, + hidden_size: int, + half_begin: int, + num_halves: int, + col_lo: int, + col_hi: int, + ): + reference = RMSNormKernel(cutlass.BFloat16, hidden_size, weight_bias=0.0) + if reference.cluster_n != 1: + raise ValueError( + f"hidden size {hidden_size} needs a CTA cluster per row in RMSNormKernel; " + "the draft glue kernel covers single-CTA rows only" + ) + self.H = hidden_size + self.weight_bias = reference.weight_bias + self.copy_bits = reference.copy_bits + self.vec_size = reference.vec_size + self.threads_per_row = reference.threads_per_row + self.num_threads = reference.num_threads + self.rows_per_block = reference.rows_per_block + self.warps_per_row = reference.warps_per_row + self.num_vec_blocks = reference.num_vec_blocks + self.cols_per_tile = reference.cols_per_tile + self.use_async_copy = reference.use_async_copy + + self.half_begin = half_begin + self.num_halves = num_halves + self.out_cols = col_hi - col_lo + self.out_col_shift = col_lo + # Stored columns [lo, hi) of each half, in that half's column coordinates. + windows = [] + for h in range(2): + lo = min(max(col_lo - h * hidden_size, 0), hidden_size) + hi = max(min(col_hi - h * hidden_size, hidden_size), lo) + windows.append((lo, hi)) + self.windows = tuple(windows) + + def _smem_size_in_bytes(self) -> int: + tile_bytes = 0 + if self.use_async_copy: + tile_bytes = self.rows_per_block * self.cols_per_tile * (cutlass.BFloat16.width // 8) + return tile_bytes + self.rows_per_block * self.warps_per_row * 4 + + @cute.jit + def __call__( + self, + mSrc0: cute.Tensor, + mIds: cute.Tensor, + mSrc1: cute.Tensor, + mW0: cute.Tensor, + mW1: cute.Tensor, + mOut: cute.Tensor, + M: Int64, + vocab: Int64, + eps: Float32, + enable_pdl: cutlass.Constexpr[bool], + stream, + ): + tv_shape, tv_stride = make_tv_layout( + self.threads_per_row, self.vec_size, self.num_vec_blocks + ) + tv_layout = cute.make_layout(tv_shape, stride=tv_stride) + tiler_1d = (1, self.cols_per_tile) + + self.kernel( + mSrc0, mIds, mSrc1, mW0, mW1, mOut, M, vocab, eps, enable_pdl, tv_layout, tiler_1d + ).launch( + grid=[cute.ceil_div(M * self.num_halves, self.rows_per_block), 1, 1], + block=[self.num_threads, 1, 1], + smem=self._smem_size_in_bytes(), + stream=stream, + use_pdl=enable_pdl, + ) + + @cute.kernel + def kernel( + self, + mSrc0: cute.Tensor, + mIds: cute.Tensor, + mSrc1: cute.Tensor, + mW0: cute.Tensor, + mW1: cute.Tensor, + mOut: cute.Tensor, + M: Int64, + vocab: Int64, + eps: Float32, + enable_pdl: cutlass.Constexpr[bool], + tv_layout: cute.Layout, + tiler_1d: cute.Shape, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + + # PDL: the ids and the hidden rows may come from the preceding kernel, so wait before any load. The + # dependent (eh_proj) may launch right away; it waits before it reads this kernel's output. + if enable_pdl: + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + + H = self.H + weight_bias = self.weight_bias + copy_bits = self.copy_bits + threads_per_row = self.threads_per_row + rows_per_block = self.rows_per_block + warps_per_row = self.warps_per_row + num_halves = self.num_halves + + lane_in_row = tidx % threads_per_row + row_in_block = tidx // threads_per_row + slot = bidx * rows_per_block + row_in_block + if cutlass.const_expr(num_halves == 1): + row = slot + is_half0 = self.half_begin == 0 + else: + row = slot // num_halves + is_half0 = slot % num_halves == 0 + row_in_bounds = row < M + + # ===== Allocate shared memory ===== + smem = cutlass.utils.SmemAllocator() + + if cutlass.const_expr(self.use_async_copy): + sX = smem.allocate_tensor( + mSrc1.element_type, + cute.make_ordered_layout((rows_per_block, self.cols_per_tile), order=(1, 0)), + byte_alignment=16, + ) + + reduction_buffer = smem.allocate_tensor( + Float32, + cute.make_layout((rows_per_block, warps_per_row)), + byte_alignment=4, + ) + mbar_ptr = None + cluster_n = 1 + + # ===== Source rows ===== + # Half 0 reads the embedding row of the token, half 1 the hidden row. + valid0 = row_in_bounds + src_row0 = row + if cutlass.const_expr(self.half_begin == 0): + # Same guard as the Embedding module: an id outside [0, vocab), such as the -1 padding of rejected draft + # tokens, gives a zero embedding row. + token = Int64(-1) + if row_in_bounds: + token = Int64(mIds[row]) + valid0 = (token >= Int64(0)) & (token < vocab) + src_row0 = cutlass.select_(valid0, token, Int64(0)) + gX0 = cute.local_tile(mSrc0, tiler_1d, (src_row0, 0)) + gX1 = cute.local_tile(mSrc1, tiler_1d, (row, 0)) + + # ===== Output: the window columns of each half ===== + # Column k of half h is stored at out column h * H + k - out_col_shift. + out_layout = cute.make_layout((M, H), stride=(self.out_cols, 1)) + mO0 = cute.make_tensor(mOut.iterator + (0 - self.out_col_shift), out_layout) + mO1 = cute.make_tensor(mOut.iterator + (H - self.out_col_shift), out_layout) + gO0 = cute.local_tile(mO0, tiler_1d, (row, 0)) + gO1 = cute.local_tile(mO1, tiler_1d, (row, 0)) + + mW0_2d = cute.prepend_ones(mW0, up_to_rank=2) + mW1_2d = cute.prepend_ones(mW1, up_to_rank=2) + + # ===== Create TiledCopy atoms ===== + copy_atom_sync = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + mSrc1.element_type, + num_bits_per_copy=copy_bits, + ) + copy_atom_store = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + mOut.element_type, + num_bits_per_copy=copy_bits, + ) + + if cutlass.const_expr(self.use_async_copy): + copy_atom_async = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(), + mSrc1.element_type, + num_bits_per_copy=copy_bits, + ) + tiled_copy_load = cute.make_tiled_copy(copy_atom_async, tv_layout, tiler_1d) + else: + tiled_copy_load = cute.make_tiled_copy(copy_atom_sync, tv_layout, tiler_1d) + + tiled_copy_W = cute.make_tiled_copy(copy_atom_sync, tv_layout, tiler_1d) + tiled_copy_store = cute.make_tiled_copy(copy_atom_store, tv_layout, tiler_1d) + + thr_copy_X = tiled_copy_load.get_slice(lane_in_row) + thr_copy_W = tiled_copy_W.get_slice(lane_in_row) + thr_copy_O = tiled_copy_store.get_slice(lane_in_row) + + # ===== Partition input ===== + tXgX0 = thr_copy_X.partition_S(gX0) + tXgX1 = thr_copy_X.partition_S(gX1) + tXrX = cute.make_fragment_like(tXgX1) + + if cutlass.const_expr(self.use_async_copy): + sX_row = cute.local_tile(sX, tiler_1d, (row_in_block, 0)) + tXsX = thr_copy_X.partition_D(sX_row) + + # ===== Partition weight (sync, separate tiled copy) ===== + tWgW0 = thr_copy_W.partition_S(mW0_2d) + tWgW1 = thr_copy_W.partition_S(mW1_2d) + tWrW = cute.make_fragment_like(tWgW0) + tXrW = thr_copy_X.retile(tWrW) + + # ===== Partition output ===== + tXgO0 = thr_copy_O.partition_D(gO0) + tXgO1 = thr_copy_O.partition_D(gO1) + tXrO = cute.make_fragment_like(tXgO0) + + # ===== Bounds checking ===== + id1d = cute.make_identity_tensor(tiler_1d) + tXcX = thr_copy_X.partition_S(id1d) + tXpX = predicate_k(tXcX, limit=H) + tWpW = predicate_k(thr_copy_W.partition_S(id1d), limit=H) + lo0, hi0 = self.windows[0] + lo1, hi1 = self.windows[1] + if cutlass.const_expr(lo0 == 0 and hi0 == H): + tXpO0 = tXpX + else: + tXpO0 = _predicate_cols(tXcX, lo0, hi0) + if cutlass.const_expr(lo1 == 0 and hi1 == H): + tXpO1 = tXpX + else: + tXpO1 = _predicate_cols(tXcX, lo1, hi1) + + # ===== Pass 1: Load input + compute sum of squares ===== + if cutlass.const_expr(self.use_async_copy): + if is_half0: + if valid0: + cute.copy(copy_atom_async, tXgX0, tXsX, pred=tXpX) + else: + tXrX.store(cute.zeros_like(tXrX, dtype=mSrc1.element_type)) + cute.autovec_copy(tXrX, tXsX) + else: + if row_in_bounds: + cute.copy(copy_atom_async, tXgX1, tXsX, pred=tXpX) + cute.arch.cp_async_commit_group() + + if is_half0: + cute.copy(copy_atom_sync, tWgW0, tWrW, pred=tWpW) + else: + cute.copy(copy_atom_sync, tWgW1, tWrW, pred=tWpW) + + cute.arch.cp_async_wait_group(0) + + cute.autovec_copy(tXsX, tXrX) + else: + tXrX.store(cute.zeros_like(tXrX, dtype=mSrc1.element_type)) + if is_half0: + if valid0: + cute.copy(copy_atom_sync, tXgX0, tXrX, pred=tXpX) + cute.copy(copy_atom_sync, tWgW0, tWrW, pred=tWpW) + else: + if row_in_bounds: + cute.copy(copy_atom_sync, tXgX1, tXrX, pred=tXpX) + cute.copy(copy_atom_sync, tWgW1, tWrW, pred=tWpW) + + x = tXrX.load().to(Float32) + x_sq = x * x + sum_sq = row_reduce_sum_multirow( + x_sq, threads_per_row, reduction_buffer, mbar_ptr, cluster_n + ) + + mean_sq = sum_sq / Float32(H) + rstd = cute.math.rsqrt(mean_sq + eps, fastmath=True) + + cute.arch.barrier() + + # ===== Pass 2: Normalize and store output ===== + # Re-load x from shared memory to relieve register pressure. + if cutlass.const_expr(self.use_async_copy): + cute.autovec_copy(tXsX, tXrX) + x = tXrX.load().to(Float32) + + w = tXrW.load().to(Float32) + y = x * rstd * (w + Float32(weight_bias)) + + tXrO.store(y.to(mOut.element_type)) + + if row_in_bounds: + if is_half0: + cute.copy(copy_atom_store, tXrO, tXgO0, pred=tXpO0) + else: + cute.copy(copy_atom_store, tXrO, tXgO1, pred=tXpO1) + + +@functools.cache +def _compile_draft_glue_kernel( + hidden_size: int, + half_begin: int, + num_halves: int, + ids_dtype: torch.dtype, + col_lo: int, + col_hi: int, + enable_pdl: bool, +): + kernel_obj = DraftGlueRMSNormKernel(hidden_size, half_begin, num_halves, col_lo, col_hi) + + sym_m = cute.sym_int(64) + src0_rows = cute.sym_int(64) if half_begin == 0 else sym_m + src0_fake = cute.runtime.make_fake_tensor( + cutlass.BFloat16, + (src0_rows, hidden_size), + (cute.sym_int64(divisibility=kernel_obj.vec_size), 1), + assumed_align=16, + ) + ids_fake = cute.runtime.make_fake_compact_tensor(_IDS_DTYPES[ids_dtype], (sym_m,)) + src1_fake = cute.runtime.make_fake_tensor( + cutlass.BFloat16, + (sym_m, hidden_size), + (cute.sym_int64(divisibility=kernel_obj.vec_size), 1), + assumed_align=16, + ) + w0_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, (hidden_size,), assumed_align=16 + ) + w1_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, (hidden_size,), assumed_align=16 + ) + out_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, (sym_m, col_hi - col_lo), stride_order=(1, 0), assumed_align=16 + ) + stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + + # Same compile options as FlashInfer's RMSNorm kernels. + return cute.compile( + kernel_obj, + src0_fake, + ids_fake, + src1_fake, + w0_fake, + w1_fake, + out_fake, + Int64(1), + Int64(1), + Float32(1e-6), + enable_pdl, + stream_fake, + options="--enable-tvm-ffi", + ) + + +def draft_glue_rmsnorm_cat( + embed_table: torch.Tensor, + input_ids: torch.Tensor, + hidden_states: torch.Tensor, + enorm_weight: torch.Tensor, + hnorm_weight: torch.Tensor, + eps: float, + col_lo: int, + col_hi: int, + enable_pdl: bool, +) -> torch.Tensor: + """Columns ``[col_lo, col_hi)`` of ``cat([rmsnorm(embed_table[input_ids]), rmsnorm(hidden_states)], -1)``. + + The embedding rows are normalized with ``enorm_weight``, the hidden rows with ``hnorm_weight``, both with + ``eps``; only the halves overlapping the window are computed. Returns the window as a contiguous + ``[M, col_hi - col_lo]`` BF16 tensor. + """ + num_rows, hidden_size = hidden_states.shape + half_begin = col_lo // hidden_size + num_halves = (col_hi - 1) // hidden_size - half_begin + 1 + # Without the embedding half the table is not read; any [M, H] tensor fills the argument. + embed_src = embed_table if half_begin == 0 else hidden_states + out = hidden_states.new_empty((num_rows, col_hi - col_lo)) + if num_rows > 0: + kernel = _compile_draft_glue_kernel( + hidden_size, half_begin, num_halves, input_ids.dtype, col_lo, col_hi, enable_pdl + ) + kernel( + embed_src, + input_ids, + hidden_states, + enorm_weight, + hnorm_weight, + out, + num_rows, + embed_src.shape[0], + eps, + ) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_glue_rmsnorm_cat.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_glue_rmsnorm_cat.py new file mode 100644 index 000000000000..45dfe48fb242 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_glue_rmsnorm_cat.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""This rank's MTP eh_proj input window -- embedding gather, enorm, hnorm, concat, column split -- in one kernel.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._draft_glue_rmsnorm import draft_glue_rmsnorm_cat as _draft_glue_rmsnorm_cat + +#: GLM-5.2's hidden size and the eh_proj input width (the concat of two hidden rows). +_HIDDEN = 6144 +#: Columns of the eh_proj input each tp16 rank holds (row-parallel eh_proj). +_TP16_WINDOW = 2 * _HIDDEN // 16 + + +def _tp16_window(rank: int) -> tuple[int, int]: + return rank * _TP16_WINDOW, (rank + 1) * _TP16_WINDOW + + +class GlmDraftGlueRmsnormCat(OpWrapper): + """Columns `[col_lo, col_hi)` of `cat([rmsnorm(embed_table[input_ids]), rmsnorm(hidden_states)], -1)`. + + The embedding rows are normalized with `enorm_weight`, the hidden rows with `hnorm_weight`, both with `eps`: + `bf16((x * rsqrt(mean(x^2) + eps)) * w)` in FP32. An id outside `[0, embed_table.shape[0])` gives a zero + embedding row, as the Embedding module does. Only the halves the window overlaps are computed and only the + window is stored. Returns a contiguous BF16 `[M, col_hi - col_lo]` tensor. + + With `enable_pdl` the kernel waits for its producer before any load and lets its dependent (eh_proj) launch + right after that wait. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # Draft step 0 at c1: the 8 verified rows of the one request (accepted-token ids, the target's last hidden + # rows). eh_proj is row-parallel over the 12288-column concat, so rank r feeds it columns + # [768 r, 768 (r + 1)): ranks 0-7 normalize only the embedding half, ranks 8-15 only the hidden half. The + # spec's `vocab` is the table's row count; the gather does not depend on it. + CELLS: tuple[Cell, ...] = ( + Cell( + why="c1 step 0, tp16 rank 0: the first embedding-half window", + spec=dict(rows=8, vocab=4096, window=_tp16_window(0), ids="valid"), + ), + Cell( + why="c1 step 0, tp16 rank 7: the last embedding-half window", + spec=dict(rows=8, vocab=4096, window=_tp16_window(7), ids="valid"), + ), + Cell( + why="c1 step 0, tp16 rank 8: the first hidden-half window", + spec=dict(rows=8, vocab=4096, window=_tp16_window(8), ids="valid"), + ), + Cell( + why="c1 step 0, tp16 rank 15: the last hidden-half window", + spec=dict(rows=8, vocab=4096, window=_tp16_window(15), ids="valid"), + ), + Cell( + why="ids outside the table (-1 padding, past the end) give zero rows", + spec=dict(rows=8, vocab=4096, window=_tp16_window(3), ids="outside"), + ), + Cell( + why="one row: a draft step that feeds the glue a single token", + spec=dict(rows=1, vocab=4096, window=_tp16_window(5), ids="valid"), + ), + ) + + note = """ + The kernel covers hidden sizes whose FlashInfer RMSNormKernel runs one CTA per row (6144 here); a window over + both halves (tp below 2, or attention DP) is supported but no target of this family reaches it. Each row is + normalized with FlashInfer's CuTe RMSNorm thread geometry, reduction and arithmetic statements; the cells + certify it against the FP32 reference. + """ + + def __call__( + self, + embed_table: torch.Tensor, + input_ids: torch.Tensor, + hidden_states: torch.Tensor, + enorm_weight: torch.Tensor, + hnorm_weight: torch.Tensor, + eps: float, + col_lo: int, + col_hi: int, + enable_pdl: bool, + ) -> torch.Tensor: + return _draft_glue_rmsnorm_cat( + embed_table, + input_ids, + hidden_states, + enorm_weight, + hnorm_weight, + eps, + col_lo, + col_hi, + enable_pdl, + ) + + def reference( + self, + embed_table: torch.Tensor, + input_ids: torch.Tensor, + hidden_states: torch.Tensor, + enorm_weight: torch.Tensor, + hnorm_weight: torch.Tensor, + eps: float, + col_lo: int, + col_hi: int, + enable_pdl: bool, + ) -> torch.Tensor: + """The unfused chain in FP32: gather with the zero-row guard, both RMSNorms, concat, column slice.""" + vocab = embed_table.shape[0] + ids = input_ids.long() + valid = (ids >= 0) & (ids < vocab) + embeds = embed_table[ids.clamp(0, vocab - 1)].float() * valid.unsqueeze(1) + + def rmsnorm(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) + return ((x * rstd) * weight.float()).to(torch.bfloat16) + + concat = torch.cat( + [rmsnorm(embeds, enorm_weight), rmsnorm(hidden_states.float(), hnorm_weight)], dim=-1 + ) + return concat[:, col_lo:col_hi].contiguous() + + def is_valid( + self, + embed_table: torch.Tensor, + input_ids: torch.Tensor, + hidden_states: torch.Tensor, + enorm_weight: torch.Tensor, + hnorm_weight: torch.Tensor, + eps: float, + col_lo: int, + col_hi: int, + enable_pdl: bool, + ) -> None: + hidden = hidden_states.shape[1] + # The predicate stores whole 8-column vectors by their first column: a window edge inside a vector moves + # columns into or out of the window. + assert 0 <= col_lo < col_hi <= 2 * hidden and col_lo % 8 == 0 and col_hi % 8 == 0, ( + "the window must be 8-column aligned inside [0, 2 * hidden)" + ) + # The kernel is compiled for 16-byte aligned rows whose stride is a multiple of 8 elements. + rows = [("hidden_states", hidden_states)] + if col_lo < hidden: + rows.append(("embed_table", embed_table)) + assert input_ids.dim() == 1 and input_ids.is_contiguous(), ( + "input_ids must be a contiguous vector" + ) + assert input_ids.shape[0] == hidden_states.shape[0], "one id per hidden row" + for name, t in rows: + assert t.dim() == 2 and t.shape[1] == hidden and t.stride(1) == 1, ( + f"{name} must be [rows, hidden] with unit column stride" + ) + assert t.stride(0) % 8 == 0 and t.data_ptr() % 16 == 0, ( + f"{name} rows must be 16-byte aligned" + ) + for name, w in (("enorm_weight", enorm_weight), ("hnorm_weight", hnorm_weight)): + assert w.shape == (hidden,) and w.is_contiguous() and w.data_ptr() % 16 == 0, ( + f"{name} must be a contiguous, 16-byte aligned [hidden] vector" + ) + + +glm_draft_glue_rmsnorm_cat = GlmDraftGlueRmsnormCat() diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py new file mode 100644 index 000000000000..ac731ef79291 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm_draft_glue_rmsnorm_cat entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec.draft_glue_rmsnorm_cat import ( # noqa: E501 + glm_draft_glue_rmsnorm_cat as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" +) + +HIDDEN = 6144 +EPS = 1e-5 + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """One RMSNorm's worth of rounding differences. + + Kernel and reference normalize the same BF16 rows and round once to BF16. They differ in the order of the 6144 + squares in the FP32 sum and in the reciprocal square root (the kernel's approximate `rsqrt`, a few FP32 ulp, the + reference's correctly rounded one): together a few 1e-7 relative on the scale factor, which moves the BF16 + rounding of an output only when it lies that close to a rounding midpoint, by one ulp of that element. + + So: 1 ulp element-wise against the row scale, 0.25 ulp of relative RMS. A wrong row, column or weight is off by + the size of the values themselves. + """ + assert_within_ulp(out, ref, element_ulp=1.0, rms_ulp=0.25) + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + rows, vocab = spec["rows"], spec["vocab"] + table = torch.randn(vocab, HIDDEN, generator=g, device="cuda", dtype=torch.bfloat16) + hidden = 3.0 * torch.randn(rows, HIDDEN, generator=g, device="cuda") + if spec["ids"] == "valid": + ids = torch.randint(0, vocab, (rows,), generator=g, device="cuda", dtype=torch.int32) + else: + # int64 ids: -1 padding, the first id past the table, valid ids, a large negative and a far one. + candidates = [-1, vocab, 5, -7, vocab + 100, 0, vocab - 1, -1] + ids = torch.tensor(candidates[:rows], dtype=torch.int64, device="cuda") + enorm = (1.0 + 0.1 * torch.randn(HIDDEN, generator=g, device="cuda")).to(torch.bfloat16) + hnorm = (1.0 + 0.1 * torch.randn(HIDDEN, generator=g, device="cuda")).to(torch.bfloat16) + lo, hi = spec["window"] + return table, ids, hidden.to(torch.bfloat16), enorm, hnorm, EPS, lo, hi + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + out = op(*args, True) + assert out.is_contiguous() and out.shape == (cell.spec["rows"], args[7] - args[6]) + _compare(out, op.reference(*args, True)) + + +def test_window_is_flashinfer_rmsnorm_bitwise() -> None: + """Where FlashInfer runs its CuTe RMSNorm, the window is that kernel's output on each half, bit for bit.""" + import flashinfer.norm + + device = torch.device("cuda") + if flashinfer.norm._use_cuda_norm(device): + pytest.skip("flashinfer.norm.rmsnorm runs its CUDA JIT kernel on this device") + for rank in (0, 15): + table, ids, hidden, enorm, hnorm, eps, lo, hi = _build( + dict(rows=8, vocab=4096, window=(768 * rank, 768 * (rank + 1)), ids="valid"), seed=rank + ) + embeds = torch.nn.functional.embedding(ids.long(), table) + enormed = flashinfer.norm.rmsnorm(embeds, enorm, eps) + hnormed = flashinfer.norm.rmsnorm(hidden, hnorm, eps) + unfused = torch.cat([enormed, hnormed], dim=-1) + out = op(table, ids, hidden, enorm, hnorm, eps, lo, hi, True) + torch.testing.assert_close(out, unfused[:, lo:hi], rtol=0, atol=0) + + +def test_guard_refuses_a_window_inside_a_vector() -> None: + """The store predicate tests each 8-column vector's first column only.""" + table, ids, hidden, enorm, hnorm, eps, _, _ = _build(op.CELLS[0].spec, seed=1) + with pytest.raises(AssertionError, match="8-column aligned"): + with validating(op): + op(table, ids, hidden, enorm, hnorm, eps, 4, 772, True) + + +def test_guard_refuses_a_misaligned_hidden_row() -> None: + """The kernel is compiled for 16-byte aligned rows.""" + table, ids, hidden, enorm, hnorm, eps, lo, hi = _build(op.CELLS[2].spec, seed=2) + storage = torch.zeros(8, HIDDEN + 8, dtype=torch.bfloat16, device="cuda") + shifted = storage[:, 1 : HIDDEN + 1] + shifted.copy_(hidden) + with pytest.raises(AssertionError, match="16-byte aligned"): + with validating(op): + op(table, ids, shifted, enorm, hnorm, eps, lo, hi, True) From 45a93c3f7629d8d2437e38eb7040338d6b7ba61e Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:12:09 -0700 Subject: [PATCH 26/53] Add the GLM-5.2 MTP draft-step kernels after the draft LM head Two collective ops replace everything between the vocabulary-sharded draft LM head and the next draft step's attention (glmMoeDsa/glmDraftSampleGlue.cu). Both take the argmax of the rank's logits shard, exchange the (index, value) pairs through a dedicated MNNVL Lamport workspace (an unwritten word is 0xFFFFFFFF) and pick the global argmax in torch.argmax's order; both write the next step's position ids, sequence and KV lengths (fill, + 1 or the step-0 rewind) and the DSA decode metadata the fused DSA decode-metadata kernel writes. - torch.ops.trtllm.glm_draft_sample_glue, one CTA: optionally the next step's eh_proj input window of cat([enorm(embed(token)), hnorm(hidden)], -1) (FlashInfer's RMSNorm arithmetic) and, after the last draft step, the next-new-tokens rows. - torch.ops.trtllm.glm_draft_ehproj_table, one 8-CTA cluster: also the next step's eh_proj output with its all-reduce and input_layernorm. Each rank publishes, with its packet, the token-table row of its candidate (eh_proj's embedding half applied to enorm(embed(t)) for the tokens of its vocabulary shard) and its FP32 share of eh_proj's hidden half; every rank writes x = bf16(row of the winner + hidden half) and input_layernorm(x). 2 to 16 ranks: at 16 a CTA's 768 columns span two ranks' 384-column shares (16 rank slots). Under PDL both read only constant norm weights before their grid-dependency wait and release their dependents right after it. The tp16 GLM target gets the entries ops/spec/draft_sample_glue and ops/spec/draft_ehproj_table, the workspace (_draft_exchange, one per op, built over the TP group with the MNNVL all-reduce workspace's communicator and handle kind) and the table's construction-time operands (_draft_token_table: the token table of the rank's shard, enorm by the draft glue kernel and the product in FP32, and the rank's hidden-half weight rows). A 4-rank matrix (_draft_step_op_matrix, started by the two collected tests) drives every cell, a c1 draft loop of six table steps and the last glue step eagerly and as CUDA-graph replays, and the guards; the same script runs 16 ranks over four nodes under an external launcher. Signed-off-by: Dongfeng Yu --- .../kernels/glmMoeDsa/glmDraftSampleGlue.cu | 998 ++++++++++++++++++ .../kernels/glmMoeDsa/glmDraftSampleGlue.h | 188 ++++ cpp/tensorrt_llm/thop/CMakeLists.txt | 3 +- .../thop/glmDraftSampleGlueOp.cpp | 366 +++++++ .../ops/spec/_draft_exchange.py | 168 +++ .../ops/spec/_draft_step_state.py | 189 ++++ .../ops/spec/_draft_token_table.py | 70 ++ .../ops/spec/draft_ehproj_table.py | 284 +++++ .../ops/spec/draft_sample_glue.py | 322 ++++++ .../_torch/custom_ops/cpp_custom_ops.py | 20 + .../spec/_draft_step_op_matrix.py | 627 +++++++++++ .../spec/_draft_step_rank_job.py | 67 ++ .../test_modeling_v2_draft_ehproj_table.py | 28 + .../test_modeling_v2_draft_sample_glue.py | 28 + .../test_modeling_v2_draft_token_table.py | 59 ++ 15 files changed, 3416 insertions(+), 1 deletion(-) create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.cu create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.h create mode 100644 cpp/tensorrt_llm/thop/glmDraftSampleGlueOp.cpp create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_exchange.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_step_state.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_token_table.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_ehproj_table.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_sample_glue.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_rank_job.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_ehproj_table.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_sample_glue.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_token_table.py diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.cu new file mode 100644 index 000000000000..ea0458c3c894 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.cu @@ -0,0 +1,998 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/common/lamportUtils.cuh" + +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ +namespace +{ + +using tensorrt_llm::common::loadPackedVolatile; +using tensorrt_llm::common::VolatilePackedLoad; + +constexpr int kBlockSize = 512; +constexpr int kNumWarps = kBlockSize / 32; +// Logit vectors (8 BF16) each thread loads before it compares them. +constexpr int kScanBatch = 8; +// FlashInfer RMSNormKernel geometry for H = 6144: thread t of a 64-thread group owns the 8-element vectors at columns +// 8 t + 512 b, b = 0..11. +constexpr int kNormThreads = 64; +constexpr int kNormVecElts = 8; +constexpr int kNormVecs = kDraftSampleGlueHidden / (kNormThreads * kNormVecElts); +constexpr int kNumNormGroups = kBlockSize / kNormThreads; +static_assert(kNormVecs * kNormThreads * kNormVecElts == kDraftSampleGlueHidden, "hidden size must be 512-aligned"); +// An unwritten word of the Lamport workspace (the MNNVL all-gather's empty word). +constexpr uint32_t kEmptyWord = 0xFFFFFFFFU; +constexpr int kMaxSlots = kDraftSampleGlueMaxRows * 32; + +struct Candidate +{ + float value; + int index; +}; + +// Total order of at::native GreaterOrNan (the order torch.max and torch.argmax use): a NaN beats every number and the +// lowest-index NaN wins; otherwise the larger value, and the lower index among equal values. +__device__ __forceinline__ bool isBetter(float v1, int i1, float v2, int i2) +{ + bool const nan1 = v1 != v1; + bool const nan2 = v2 != v2; + if (nan1 || nan2) + { + return nan1 && (!nan2 || i1 < i2); + } + return v1 > v2 || (v1 == v2 && i1 < i2); +} + +__device__ __forceinline__ void offer(Candidate& best, float value, int index) +{ + if (isBetter(value, index, best.value, best.index)) + { + best.value = value; + best.index = index; + } +} + +__device__ __forceinline__ Candidate warpArgmax(Candidate c) +{ +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + { + float const value = __shfl_xor_sync(0xFFFFFFFFU, c.value, offset); + int const index = __shfl_xor_sync(0xFFFFFFFFU, c.index, offset); + offer(c, value, index); + } + return c; +} + +// BF16 element 2k (low half) and 2k + 1 (high half) of a 32-bit word, widened exactly. +__device__ __forceinline__ float bf16Low(uint32_t word) +{ + return __uint_as_float(word << 16); +} + +__device__ __forceinline__ float bf16High(uint32_t word) +{ + return __uint_as_float(word & 0xFFFF0000U); +} + +__device__ __forceinline__ void offerVector(Candidate& best, uint4 const& v, int col) +{ + offer(best, bf16Low(v.x), col); + offer(best, bf16High(v.x), col + 1); + offer(best, bf16Low(v.y), col + 2); + offer(best, bf16High(v.y), col + 3); + offer(best, bf16Low(v.z), col + 4); + offer(best, bf16High(v.z), col + 5); + offer(best, bf16Low(v.w), col + 6); + offer(best, bf16High(v.w), col + 7); +} + +// This thread's best element of a logits row (vocabShard / 8 vectors of 8 BF16). +__device__ Candidate scanRow(uint4 const* row, int numVecs) +{ + Candidate best{-INFINITY, INT_MAX}; + for (int base = static_cast(threadIdx.x); base < numVecs; base += kScanBatch * kBlockSize) + { + uint4 vecs[kScanBatch]; +#pragma unroll + for (int k = 0; k < kScanBatch; ++k) + { + int const v = base + k * kBlockSize; + if (v < numVecs) + { + vecs[k] = row[v]; + } + } +#pragma unroll + for (int k = 0; k < kScanBatch; ++k) + { + int const v = base + k * kBlockSize; + if (v < numVecs) + { + offerVector(best, vecs[k], v * kNormVecElts); + } + } + } + return best; +} + +__device__ __forceinline__ bool isEmptyPacket(VolatilePackedLoad const& packet) +{ + return packet.words[0] == kEmptyWord || packet.words[1] == kEmptyWord || packet.words[2] == kEmptyWord + || packet.words[3] == kEmptyWord; +} + +__device__ __forceinline__ void namedBarrier(int id, int numThreads) +{ + asm volatile("bar.sync %0, %1;" ::"r"(id), "r"(numThreads) : "memory"); +} + +__device__ __forceinline__ float rsqrtApprox(float x) +{ + float y; + asm("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ uint32_t packBf16(float lo, float hi) +{ + __nv_bfloat162 const v = __floats2bfloat162_rn(lo, hi); + return *reinterpret_cast(&v); +} + +// (x * rstd) * (w + 0) of 8 elements, rounded to BF16. +__device__ __forceinline__ uint4 normalizeVector(uint4 const& x, uint4 const& w, float rstd) +{ + auto const scaled = [rstd](float xv, float wv) { return __fmul_rn(__fmul_rn(xv, rstd), __fadd_rn(wv, 0.F)); }; + uint4 y; + y.x = packBf16(scaled(bf16Low(x.x), bf16Low(w.x)), scaled(bf16High(x.x), bf16High(w.x))); + y.y = packBf16(scaled(bf16Low(x.y), bf16Low(w.y)), scaled(bf16High(x.y), bf16High(w.y))); + y.z = packBf16(scaled(bf16Low(x.z), bf16Low(w.z)), scaled(bf16High(x.z), bf16High(w.z))); + y.w = packBf16(scaled(bf16Low(x.w), bf16Low(w.w)), scaled(bf16High(x.w), bf16High(w.w))); + return y; +} + +__device__ __forceinline__ float sumOfSquares(uint4 const& v, float acc) +{ + uint32_t const words[4] = {v.x, v.y, v.z, v.w}; +#pragma unroll + for (int k = 0; k < 4; ++k) + { + float const lo = bf16Low(words[k]); + float const hi = bf16High(words[k]); + acc = __fadd_rn(acc, __fmul_rn(lo, lo)); + acc = __fadd_rn(acc, __fmul_rn(hi, hi)); + } + return acc; +} + +// RMSNorm of one row (src, or a zero row when null) by one 64-thread group, with the arithmetic of FlashInfer's +// RMSNormKernel: each thread sums the squares of its 96 elements in column order, the warp adds the partial sums in a +// butterfly (offsets 1, 2, ..., 16), the row sum is the first warp's total plus the second's; then +// rsqrt.approx(sum / H + eps) and (x * rstd) * (w + 0) in FP32, rounded to BF16. Columns [lo, hi) of the row (half +// coordinates) are stored at out[col - lo]. +__device__ void normRowWindow(uint4 const* src, __nv_bfloat16 const* weights, float eps, int lo, int hi, + __nv_bfloat16* out, int group, float (*groupSums)[2]) +{ + int const t = static_cast(threadIdx.x) % kNormThreads; + int const warpInGroup = t / 32; + int const lane = t % 32; + uint4 x[kNormVecs]; +#pragma unroll + for (int b = 0; b < kNormVecs; ++b) + { + x[b] = src != nullptr ? src[t + kNormThreads * b] : make_uint4(0U, 0U, 0U, 0U); + } + float acc = 0.F; +#pragma unroll + for (int b = 0; b < kNormVecs; ++b) + { + acc = sumOfSquares(x[b], acc); + } +#pragma unroll + for (int offset = 1; offset < 32; offset <<= 1) + { + acc = __fadd_rn(acc, __shfl_xor_sync(0xFFFFFFFFU, acc, offset)); + } + if (lane == 0) + { + groupSums[group][warpInGroup] = acc; + } + namedBarrier(1 + group, kNormThreads); + float const sum = __fadd_rn(groupSums[group][0], groupSums[group][1]); + float const rstd = rsqrtApprox(__fadd_rn(__fdiv_rn(sum, static_cast(kDraftSampleGlueHidden)), eps)); +#pragma unroll + for (int b = 0; b < kNormVecs; ++b) + { + int const col = kNormVecElts * t + kNormThreads * kNormVecElts * b; + if (col >= lo && col < hi) + { + uint4 const w = *reinterpret_cast(weights + col); + *reinterpret_cast(out + (col - lo)) = normalizeVector(x[b], w, rstd); + } + } + // The group's sums are rewritten by its next row. + namedBarrier(1 + group, kNormThreads); +} + +// Inclusive warp prefix sum. +__device__ __forceinline__ int64_t warpInclusiveSum(int64_t value) +{ + int const lane = static_cast(threadIdx.x) % 32; +#pragma unroll + for (int offset = 1; offset < 32; offset <<= 1) + { + int64_t const other = __shfl_up_sync(0xFFFFFFFFU, value, offset); + if (lane >= offset) + { + value += other; + } + } + return value; +} + +// Position ids, sequence and KV lengths of the next step, then the DSA decode metadata of fused_dsa_decode_metadata +// (lane r handles request r). One warp. +__device__ void advanceStep(DraftSampleGlueParams const& p) +{ + int const lane = static_cast(threadIdx.x) % 32; + bool const active = lane < p.numRows; + int32_t seq = 0; + int32_t kv = 0; + if (active) + { + if (p.positionsOut != nullptr) + { + int64_t const src = p.gatherIds != nullptr ? p.gatherIds[lane] : lane; + if (p.positions64) + { + static_cast(p.positionsOut)[lane] = static_cast(p.positionsIn)[src] + 1; + } + else + { + static_cast(p.positionsOut)[lane] = static_cast(p.positionsIn)[src] + 1; + } + } + if (p.seqLens != nullptr) + { + seq = p.fillSeqLens ? 1 : p.seqLens[lane]; + if (p.fillSeqLens) + { + p.seqLens[lane] = 1; + } + } + if (p.kvLens != nullptr) + { + kv = p.kvLens[lane]; + if (p.kvUpdate == DraftSampleGlueKvUpdate::kIncrement) + { + kv += 1; + p.kvLens[lane] = kv; + } + else if (p.kvUpdate == DraftSampleGlueKvUpdate::kRewind) + { + kv -= p.draftLen - p.kvRewindAccepted[lane]; + p.kvLens[lane] = kv; + } + } + } + if (p.reqIdxPerToken == nullptr) + { + return; + } + int64_t const seq64 = active ? seq : 0; + int64_t const kv64 = active ? kv : 0; + int64_t const seqIncl = warpInclusiveSum(seq64); + int64_t const kvIncl = warpInclusiveSum(kv64); + int64_t const cachedIncl = warpInclusiveSum(kv64 - seq64); + if (lane == 0) + { + p.genKvIndptr[0] = 0; + p.genCachedTokenIndptr[0] = 0; + } + if (!active) + { + return; + } + p.genKvIndptr[1 + lane] = kvIncl; + p.genCachedTokenIndptr[1 + lane] = cachedIncl; + int64_t const flatStart = seqIncl - seq64; + int64_t const startPos = kv64 - seq64; + int64_t const tokens = seq64 < p.maxQueryLanes ? seq64 : p.maxQueryLanes; + for (int64_t j = 0; j < tokens; ++j) + { + int64_t const g = flatStart + j; + if (g >= p.numTokens) + { + break; + } + int64_t const globalPos = startPos + j; + // Truncating division and a non-negative remainder, then the clamp, as the Triton kernel computes them. + int64_t blk = globalPos / p.tokensPerBlock; + int64_t pos = globalPos % p.tokensPerBlock; + pos = pos < 0 ? pos + p.tokensPerBlock : pos; + blk = blk < 0 ? 0 : (blk > p.maxBlocks - 1 ? p.maxBlocks - 1 : blk); + int64_t const blockId = p.blockOffsets[lane * p.blockOffsetsStride0 + blk * p.blockOffsetsStride1]; + p.reqIdxPerToken[g] = lane; + p.slotMappingFp8[g] = blockId * p.blockStride + pos * p.dataBytesPerToken; + p.slotMappingScale[g] = blockId * p.blockStride + p.scaleBaseOffset + pos * p.scaleSize; + } +} + +__global__ void __launch_bounds__(kBlockSize) draftSampleGlueKernel(DraftSampleGlueParams const p) +{ + __shared__ __align__(16) __nv_bfloat16 sWeights[2][kDraftSampleGlueHidden]; + __shared__ Candidate sCandidates[kDraftSampleGlueMaxRows][kNumWarps]; + __shared__ Candidate sSlots[kMaxSlots]; + __shared__ int sTokens[kDraftSampleGlueMaxRows]; + __shared__ float sGroupSums[kNumNormGroups][2]; + __shared__ uint4 sFlags[2]; + + int const tid = static_cast(threadIdx.x); + int const warp = tid / 32; + int const lane = tid % 32; + bool const glue = p.glueOut != nullptr; + // The halves of the concatenated row that the eh_proj input window overlaps. + bool const embedHalf = glue && p.colLo < kDraftSampleGlueHidden; + bool const hiddenHalf = glue && p.colHi > kDraftSampleGlueHidden; + + // The norm weights are constants: copy the window columns before the grid dependency wait. + if (glue) + { + for (int h = 0; h < 2; ++h) + { + if ((h == 0 && !embedHalf) || (h == 1 && !hiddenHalf)) + { + continue; + } + int const lo = max(p.colLo - h * kDraftSampleGlueHidden, 0); + int const hi = min(p.colHi - h * kDraftSampleGlueHidden, kDraftSampleGlueHidden); + auto const* src = reinterpret_cast(h == 0 ? p.enormWeight : p.hnormWeight); + auto* dst = reinterpret_cast(sWeights[h]); + for (int v = lo / kNormVecElts + tid; v < hi / kNormVecElts; v += kBlockSize) + { + dst[v] = src[v]; + } + } + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + // The logits come from the preceding kernel (the draft LM head). Dependents read the outputs only after their + // own wait. + cudaGridDependencySynchronize(); + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + if (tid == 0) + { + sFlags[0] = reinterpret_cast(p.exchangeFlags)[0]; + sFlags[1] = reinterpret_cast(p.exchangeFlags)[1]; + } + + // Local argmax of every row: all threads scan, each warp reduces, warp 0 combines below. + int const numVecs = p.vocabShard / kNormVecElts; + for (int row = 0; row < p.numRows; ++row) + { + Candidate const c + = warpArgmax(scanRow(reinterpret_cast(p.logits + row * p.logitsStride), numVecs)); + if (lane == 0) + { + sCandidates[row][warp] = c; + } + } + __syncthreads(); + + uint4 const flags = sFlags[0]; + uint32_t const curIndex = flags.x; + uint32_t const dirtyIndex = flags.y; + uint32_t const bytesPerBuffer = flags.z; + uint32_t const dirtyStages = flags.w; + int const numSlots = p.numRows * p.tpSize; + auto* exchangeLocal = static_cast(p.exchangeLocal); + + if (warp == 0) + { + // Publish this rank's pair of every row through the multicast address: slot [row][rank] of the current buffer. + auto* curMcast = reinterpret_cast(static_cast(p.exchangeMcast) + curIndex * bytesPerBuffer); + for (int row = 0; row < p.numRows; ++row) + { + Candidate c = lane < kNumWarps ? sCandidates[row][lane] : Candidate{-INFINITY, INT_MAX}; + c = warpArgmax(c); + if (lane == 0) + { + int const globalIndex = c.index + p.rank * p.vocabShard; + curMcast[row * p.tpSize + p.rank] + = make_uint4(static_cast(globalIndex), __float_as_uint(c.value), 0U, 0U); + } + } + // Poll the local copy until every rank's pair has landed. + auto const* curLocal = exchangeLocal + curIndex * bytesPerBuffer; + for (int s = lane; s < numSlots; s += 32) + { + VolatilePackedLoad packet; + do + { + packet = loadPackedVolatile(curLocal + s * kDraftSampleGluePacketBytes); + } while (isEmptyPacket(packet)); + sSlots[s] = Candidate{__uint_as_float(packet.words[1]), static_cast(packet.words[0])}; + } + __syncwarp(); + // Global argmax per row: the ranks' pairs in rank order (the lower rank holds the lower indices). + if (lane < p.numRows) + { + Candidate best = sSlots[lane * p.tpSize]; + for (int r = 1; r < p.tpSize; ++r) + { + offer(best, sSlots[lane * p.tpSize + r].value, sSlots[lane * p.tpSize + r].index); + } + sTokens[lane] = best.index; + p.draftTokens[lane * p.numDraftTokens + p.draftStep] = best.index; + } + } + else if (warp == 1) + { + // Empty the buffer of the previous call; the peers write it again two calls later. + if (dirtyStages > 0) + { + auto* dirty = reinterpret_cast(exchangeLocal + dirtyIndex * bytesPerBuffer); + int const dirtyPackets = static_cast(sFlags[1].x / kDraftSampleGluePacketBytes); + for (int i = lane; i < dirtyPackets; i += 32) + { + dirty[i] = make_uint4(kEmptyWord, kEmptyWord, kEmptyWord, kEmptyWord); + } + } + advanceStep(p); + } + else if (hiddenHalf) + { + // hnorm of the hidden rows does not need the tokens: groups 1..7 (warps 2..15) run it during the poll. + int const group = warp / 2; + int const lo = max(p.colLo - kDraftSampleGlueHidden, 0); + int const hi = min(p.colHi - kDraftSampleGlueHidden, kDraftSampleGlueHidden); + int const outShift = kDraftSampleGlueHidden + lo - p.colLo; + for (int row = group - 1; row < p.numRows; row += kNumNormGroups - 1) + { + normRowWindow(reinterpret_cast(p.hidden + row * p.hiddenStride), sWeights[1], p.eps, lo, hi, + p.glueOut + row * p.glueOutStride + outShift, group, sGroupSums); + } + } + __syncthreads(); + + if (embedHalf) + { + // enorm of the token embeddings: an id outside [0, vocabSize) gives a zero row, as the Embedding module does. + int const group = warp / 2; + int const lo = p.colLo; + int const hi = min(p.colHi, kDraftSampleGlueHidden); + for (int row = group; row < p.numRows; row += kNumNormGroups) + { + int64_t const token = sTokens[row]; + bool const valid = token >= 0 && token < p.vocabSize; + normRowWindow(valid ? reinterpret_cast(p.embedTable + token * kDraftSampleGlueHidden) : nullptr, + sWeights[0], p.eps, lo, hi, p.glueOut + row * p.glueOutStride, group, sGroupSums); + } + } + + if (warp == 1 && p.nextNewTokens != nullptr && lane < p.numRows) + { + // Output rows after the last draft step: the last accepted token, then the draft tokens of every step. + int32_t* out = p.nextNewTokens + lane * (p.numDraftTokens + 1); + out[0] = p.acceptedTokens[p.batchIndices[lane] * p.acceptedTokensStride + p.numAccepted[lane] - 1]; + for (int j = 0; j < p.numDraftTokens; ++j) + { + out[1 + j] = j == p.draftStep ? sTokens[lane] : p.draftTokens[lane * p.numDraftTokens + j]; + } + } + + if (tid == 0) + { + // Rotate the Lamport flags as LamportFlags::update does: the current buffer becomes the dirty one. + auto* flagPtr = reinterpret_cast(p.exchangeFlags); + flagPtr[0] = make_uint4((curIndex + 1) % 3, curIndex, bytesPerBuffer, 1U); + flagPtr[1] = make_uint4(static_cast(numSlots * kDraftSampleGluePacketBytes), 0U, 0U, 0U); + p.exchangeFlags[8] = 0U; + } +} + +// The draft step with the next step's eh_proj, all-reduce and input_layernorm (draftEhprojTableKernel). The cluster +// covers the hidden row: thread t < kTableSliceVecs of CTA c owns columns [c * kTableSliceElts + 8 t, + 8). +constexpr int kTableCtas = kDraftEhprojTableClusterSize; +constexpr int kTableSliceElts = kDraftSampleGlueHidden / kTableCtas; +constexpr int kTableSliceVecs = kTableSliceElts / kNormVecElts; +constexpr int kTableSliceBytes = kTableSliceElts * static_cast(sizeof(__nv_bfloat16)); +constexpr int kTableVecWarps = kTableSliceVecs / 32; +// After the vector warps: a warp that polls the (index, value) packets and scans, then a warp that updates the next +// step's state (CTA 0) or empties the previous call's buffer (the other CTAs) beside the exchange. +constexpr int kTableExchangeWarp = kTableVecWarps; +constexpr int kTableStateWarp = kTableVecWarps + 1; +constexpr int kTableMainThreads = kTableStateWarp * 32; +constexpr int kTableThreads = kTableMainThreads + 32; +constexpr int kTableMainWarps = kTableMainThreads / 32; +// Rank slots of the kernel (its template argument): eight up to TP8, kDraftEhprojTableMaxRanks above. +constexpr int kTableTp8RankSlots = kTableCtas; +// Logit vectors each scanning thread loads before it compares them. +constexpr int kTableScanBatch = 4; +// Slot layout: the packet, the BF16 token-table row, then the rank's FP32 hidden-half outputs. +constexpr int kTableRowOffset = kDraftSampleGluePacketBytes; +constexpr int kTableHiddenOffset = kTableRowOffset + 2 * kDraftSampleGlueHidden; +// Named barrier of the main warps (the state warp runs on its own). +constexpr int kTableMainBarrier = 1; +static_assert(kTableSliceVecs % 32 == 0, "the slice must be whole warps of vectors"); +static_assert(kTableSliceElts * kTableCtas == kDraftSampleGlueHidden, "the cluster must cover the row"); + +__device__ __forceinline__ uint4 sanitizeWords(uint4 v) +{ + // A word with every bit set is the empty word; such a payload word (two NaN halves) is sent as another NaN. + v.x = v.x == kEmptyWord ? kEmptyWord - 1U : v.x; + v.y = v.y == kEmptyWord ? kEmptyWord - 1U : v.y; + v.z = v.z == kEmptyWord ? kEmptyWord - 1U : v.z; + v.w = v.w == kEmptyWord ? kEmptyWord - 1U : v.w; + return v; +} + +__device__ __forceinline__ float warpSum(float v) +{ +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + { + v = __fadd_rn(v, __shfl_xor_sync(0xFFFFFFFFU, v, offset)); + } + return v; +} + +__device__ __forceinline__ void clusterArriveRelaxed() +{ + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); +} + +__device__ __forceinline__ void clusterArrive() +{ + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); +} + +__device__ __forceinline__ void clusterWait() +{ + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); +} + +template +__global__ void __launch_bounds__(kTableThreads) draftEhprojTableKernel(DraftEhprojTableParams const p) +{ + // sCtaBest is read by the other CTAs of the cluster after the first cluster barrier; sClusterSums is written by + // every CTA (slot = the writer's rank in the cluster) before the second one and read locally after it. The rest + // is shared by the main warps through the named barrier. + __shared__ Candidate sWarpBest[kTableMainWarps]; + __shared__ float sWarpSums[kTableMainWarps]; + __shared__ Candidate sCtaBest[kDraftSampleGlueMaxRows]; + __shared__ Candidate sRankBest[kDraftSampleGlueMaxRows]; + __shared__ float sClusterSums[kDraftSampleGlueMaxRows][kTableCtas]; + __shared__ float sRstd[kDraftSampleGlueMaxRows]; + __shared__ Candidate sPairs[kDraftSampleGlueMaxRows][kMaxRanks]; + __shared__ __align__(16) __nv_bfloat16 sX[kDraftSampleGlueMaxRows][kTableSliceElts]; + __shared__ uint4 sFlags; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + DraftSampleGlueParams const& s = p.step; + int const tid = static_cast(threadIdx.x); + int const warp = tid / 32; + int const lane = tid % 32; + int const cta = static_cast(cluster.block_rank()); + int const tp = s.tpSize; + int const colBase = cta * kTableSliceElts; + bool const vecThread = tid < kTableSliceVecs; + // The rank whose share of the hidden-half outputs (H / tp columns) holds this vector thread's eight columns, and + // their offset in that share: a CTA's columns lie in one share up to TP8 and span two shares at TP16. + int const shareCols = kDraftSampleGlueHidden / tp; + int const vecCol = colBase + kNormVecElts * tid; + int const hiddenOwner = vecCol / shareCols; + int const shareOffset = vecCol - hiddenOwner * shareCols; + int const slotBytes = draftEhprojTableSlotBytes(tp); + + // The norm weight is a constant: load it before the grid dependency wait. + uint4 normWeight = make_uint4(0U, 0U, 0U, 0U); + if (vecThread) + { + normWeight = reinterpret_cast(p.normWeight + colBase)[tid]; + } + + // The logits come from the preceding kernel (the draft LM head), the hidden-half outputs from a kernel this launch + // depends on through an event. Dependents read the outputs only after their own wait. + cudaGridDependencySynchronize(); + cudaTriggerProgrammaticLaunchCompletion(); + + if (warp == kTableStateWarp) + { + // Arrive for the first cluster barrier before the state work, which nothing in the cluster waits for. + clusterArriveRelaxed(); + if (cta == 0) + { + advanceStep(s); + } + else + { + // Empty this CTA's share of the previous call's buffer; the peers write it again two calls later. + uint4 const flags = reinterpret_cast(s.exchangeFlags)[0]; + uint4 const sizes = reinterpret_cast(s.exchangeFlags)[1]; + if (flags.w > 0) + { + auto* dirty = reinterpret_cast( + static_cast(s.exchangeLocal) + static_cast(flags.y) * flags.z); + int const vectors = static_cast(sizes.x / sizeof(uint4)); + for (int i = (cta - 1) * 32 + lane; i < vectors; i += (kTableCtas - 1) * 32) + { + dirty[i] = make_uint4(kEmptyWord, kEmptyWord, kEmptyWord, kEmptyWord); + } + } + } + clusterWait(); + // The second (last) cluster barrier. + clusterArrive(); + clusterWait(); + return; + } + + if (tid == 0) + { + sFlags = reinterpret_cast(s.exchangeFlags)[0]; + } + + // Per row: this CTA's best logit (the CTAs scan consecutive chunks of the shard). + int const numVecs = s.vocabShard / kNormVecElts; + int const chunk = (numVecs + kTableCtas - 1) / kTableCtas; + int const vecBegin = min(cta * chunk, numVecs); + int const vecEnd = min(vecBegin + chunk, numVecs); + for (int row = 0; row < s.numRows; ++row) + { + auto const* logitsRow = reinterpret_cast(s.logits + row * s.logitsStride); + Candidate best{-INFINITY, INT_MAX}; + for (int base = vecBegin + tid; base < vecEnd; base += kTableScanBatch * kTableMainThreads) + { + uint4 vecs[kTableScanBatch]; +#pragma unroll + for (int k = 0; k < kTableScanBatch; ++k) + { + int const v = base + k * kTableMainThreads; + if (v < vecEnd) + { + vecs[k] = logitsRow[v]; + } + } +#pragma unroll + for (int k = 0; k < kTableScanBatch; ++k) + { + int const v = base + k * kTableMainThreads; + if (v < vecEnd) + { + offerVector(best, vecs[k], v * kNormVecElts); + } + } + } + best = warpArgmax(best); + if (lane == 0) + { + sWarpBest[warp] = best; + } + namedBarrier(kTableMainBarrier, kTableMainThreads); + Candidate ctaBest = sWarpBest[0]; + for (int w = 1; w < kTableMainWarps; ++w) + { + offer(ctaBest, sWarpBest[w].value, sWarpBest[w].index); + } + if (tid == 0) + { + sCtaBest[row] = ctaBest; + } + // The winner's token-table row is one of the cluster's candidates: pull this CTA's candidate row into L2 while + // the cluster combines them. + if (vecThread && ctaBest.index >= 0 && ctaBest.index < s.vocabShard) + { + auto const* row = p.tokenTable + static_cast(ctaBest.index) * kDraftSampleGlueHidden; + char const* line = reinterpret_cast(row) + tid * 128; + asm volatile("prefetch.global.L2 [%0];" ::"l"(line) : "memory"); + } + // sWarpBest is rewritten for the next row only after this barrier. + namedBarrier(kTableMainBarrier, kTableMainThreads); + } + + uint4 const flags = sFlags; + uint32_t const curIndex = flags.x; + uint32_t const bytesPerBuffer = flags.z; + char* const curMcast = static_cast(s.exchangeMcast) + static_cast(curIndex) * bytesPerBuffer; + char const* const curLocal + = static_cast(s.exchangeLocal) + static_cast(curIndex) * bytesPerBuffer; + + // This rank's hidden-half outputs do not depend on the exchange: publish them first. + if (vecThread && hiddenOwner == s.rank) + { + for (int row = 0; row < s.numRows; ++row) + { + auto const* h = reinterpret_cast(p.hiddenHalf + row * p.hiddenHalfStride + shareOffset); + char* const slot = curMcast + static_cast(row * tp + s.rank) * slotBytes; + auto* dst = reinterpret_cast(slot + kTableHiddenOffset + shareOffset * sizeof(float)); + uint4 const h0 = h[0]; + uint4 const h1 = h[1]; + dst[0] = sanitizeWords(h0); + dst[1] = sanitizeWords(h1); + } + } + + clusterArrive(); + clusterWait(); + + // This rank's candidate of every row: the CTAs' candidates in CTA order. These are the last reads of another CTA's + // shared memory; every CTA makes them before it arrives at the second cluster barrier. + if (tid < s.numRows) + { + Candidate rankBest = *cluster.map_shared_rank(&sCtaBest[tid], 0); + for (int c = 1; c < kTableCtas; ++c) + { + Candidate const other = *cluster.map_shared_rank(&sCtaBest[tid], c); + offer(rankBest, other.value, other.index); + } + sRankBest[tid] = rankBest; + } + namedBarrier(kTableMainBarrier, kTableMainThreads); + + for (int row = 0; row < s.numRows; ++row) + { + Candidate const rankBest = sRankBest[row]; + char* const slot = curMcast + static_cast(row * tp + s.rank) * slotBytes; + if (vecThread) + { + // This CTA's columns of the token-table row of the rank's candidate. + uint4 const q = reinterpret_cast( + p.tokenTable + static_cast(rankBest.index) * kDraftSampleGlueHidden + colBase)[tid]; + reinterpret_cast(slot + kTableRowOffset + cta * kTableSliceBytes)[tid] = sanitizeWords(q); + } + else if (cta == 0 && tid == kTableExchangeWarp * 32) + { + int const globalIndex = rankBest.index + s.rank * s.vocabShard; + *reinterpret_cast(slot) + = make_uint4(static_cast(globalIndex), __float_as_uint(rankBest.value), 0U, 0U); + } + + // Poll the local copy: the ranks' rows of this CTA's columns and the two vectors of hidden-half outputs that + // cover them (vector threads), the ranks' packets (exchange warp). Every round has all its loads in flight. + char const* const rowLocal = curLocal + static_cast(row * tp) * slotBytes; + VolatilePackedLoad rows[kMaxRanks]; + VolatilePackedLoad hidden0; + VolatilePackedLoad hidden1; + if (vecThread) + { + char const* const rowSlice = rowLocal + kTableRowOffset + cta * kTableSliceBytes + tid * sizeof(uint4); + char const* const hiddenPtr = rowLocal + static_cast(hiddenOwner) * slotBytes + kTableHiddenOffset + + shareOffset * sizeof(float); + bool empty = true; + while (empty) + { +#pragma unroll + for (int r = 0; r < kMaxRanks; ++r) + { + if (r < tp) + { + rows[r] = loadPackedVolatile(rowSlice + static_cast(r) * slotBytes); + } + } + hidden0 = loadPackedVolatile(hiddenPtr); + hidden1 = loadPackedVolatile(hiddenPtr + sizeof(uint4)); + empty = isEmptyPacket(hidden0) || isEmptyPacket(hidden1); +#pragma unroll + for (int r = 0; r < kMaxRanks; ++r) + { + if (r < tp) + { + empty = empty || isEmptyPacket(rows[r]); + } + } + } + } + else if (warp == kTableExchangeWarp && lane < tp) + { + VolatilePackedLoad packet; + do + { + packet = loadPackedVolatile(rowLocal + static_cast(lane) * slotBytes); + } while (isEmptyPacket(packet)); + sPairs[row][lane] = Candidate{__uint_as_float(packet.words[1]), static_cast(packet.words[0])}; + } + namedBarrier(kTableMainBarrier, kTableMainThreads); + + // The global argmax: the ranks' packets in rank order (the lower rank holds the lower indices). + int winner = 0; + Candidate best = sPairs[row][0]; + for (int r = 1; r < tp; ++r) + { + if (isBetter(sPairs[row][r].value, sPairs[row][r].index, best.value, best.index)) + { + best = sPairs[row][r]; + winner = r; + } + } + if (cta == 0 && tid == 0) + { + s.draftTokens[row * s.numDraftTokens + s.draftStep] = best.index; + } + + float sumSq = 0.F; + if (vecThread) + { + uint32_t qw[4] = {rows[0].words[0], rows[0].words[1], rows[0].words[2], rows[0].words[3]}; +#pragma unroll + for (int r = 1; r < kMaxRanks; ++r) + { +#pragma unroll + for (int k = 0; k < 4; ++k) + { + qw[k] = r == winner ? rows[r].words[k] : qw[k]; + } + } + // x = token-table row + hidden-half output: one FP32 add and one BF16 rounding per element. + uint32_t const hw[8] = {hidden0.words[0], hidden0.words[1], hidden0.words[2], hidden0.words[3], + hidden1.words[0], hidden1.words[1], hidden1.words[2], hidden1.words[3]}; + uint32_t xw[4]; +#pragma unroll + for (int k = 0; k < 4; ++k) + { + xw[k] = packBf16(__fadd_rn(bf16Low(qw[k]), __uint_as_float(hw[2 * k])), + __fadd_rn(bf16High(qw[k]), __uint_as_float(hw[2 * k + 1]))); + } + uint4 const x = make_uint4(xw[0], xw[1], xw[2], xw[3]); + reinterpret_cast(sX[row])[tid] = x; + reinterpret_cast(p.residualOut + row * p.outStride + colBase)[tid] = x; + sumSq = sumOfSquares(x, 0.F); + } + sumSq = warpSum(sumSq); + if (lane == 0) + { + sWarpSums[warp] = sumSq; + } + namedBarrier(kTableMainBarrier, kTableMainThreads); + if (tid < kTableCtas) + { + float sum = sWarpSums[0]; + for (int w = 1; w < kTableVecWarps; ++w) + { + sum = __fadd_rn(sum, sWarpSums[w]); + } + *cluster.map_shared_rank(&sClusterSums[row][cta], tid) = sum; + } + // sPairs and sWarpSums are rewritten for the next row only after this barrier. + namedBarrier(kTableMainBarrier, kTableMainThreads); + } + + // Every CTA has sent its sums of squares and read the others' candidates: after this barrier no CTA touches + // another's shared memory. + clusterArrive(); + clusterWait(); + + // input_layernorm: the row's sum of squares over the cluster in CTA order. + if (tid < s.numRows) + { + float sum = sClusterSums[tid][0]; + for (int c = 1; c < kTableCtas; ++c) + { + sum = __fadd_rn(sum, sClusterSums[tid][c]); + } + sRstd[tid] = rsqrtApprox(__fadd_rn(__fdiv_rn(sum, static_cast(kDraftSampleGlueHidden)), p.normEps)); + } + namedBarrier(kTableMainBarrier, kTableMainThreads); + if (vecThread) + { + for (int row = 0; row < s.numRows; ++row) + { + uint4 const x = reinterpret_cast(sX[row])[tid]; + reinterpret_cast(p.normedOut + row * p.outStride + colBase)[tid] + = normalizeVector(x, normWeight, sRstd[row]); + } + } + + if (cta == 0 && tid == 0) + { + // Every CTA read the flags before the first cluster barrier, every state warp before the second: rotate them + // as LamportFlags::update does. + auto* flagPtr = reinterpret_cast(s.exchangeFlags); + flagPtr[0] = make_uint4((curIndex + 1) % 3, curIndex, bytesPerBuffer, 1U); + flagPtr[1] = make_uint4(static_cast(s.numRows * tp * slotBytes), 0U, 0U, 0U); + s.exchangeFlags[8] = 0U; + } +#endif +} + +} // namespace + +void invokeDraftSampleGlue(DraftSampleGlueParams const& params, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(params.numRows > 0 && params.numRows <= kDraftSampleGlueMaxRows, + "draft sample glue serves 1 to %d rows, got %d", kDraftSampleGlueMaxRows, params.numRows); + TLLM_CHECK_WITH_INFO(params.tpSize >= 1 && params.numRows * params.tpSize <= kDraftSampleGlueMaxRows * 32, + "draft sample glue: %d rows x %d ranks exceed the exchange slots", params.numRows, params.tpSize); + TLLM_CHECK_WITH_INFO(params.vocabShard > 0 && params.vocabShard % 8 == 0, + "draft sample glue needs a vocabulary shard of a multiple of 8 columns, got %d", params.vocabShard); + TLLM_CHECK_WITH_INFO(params.glueOut == nullptr + || (params.colLo >= 0 && params.colLo < params.colHi && params.colHi <= 2 * kDraftSampleGlueHidden + && params.colLo % 8 == 0 && params.colHi % 8 == 0), + "draft sample glue: invalid eh_proj input window [%d, %d)", params.colLo, params.colHi); + TLLM_CHECK_WITH_INFO(params.glueOut == nullptr || params.colLo >= kDraftSampleGlueHidden + || params.embedTable != nullptr, + "draft sample glue: the window needs the embedding table"); + TLLM_CHECK_WITH_INFO(params.glueOut == nullptr || params.colHi <= kDraftSampleGlueHidden || params.hidden != nullptr, + "draft sample glue: the window needs the hidden rows"); + TLLM_CHECK_WITH_INFO(params.reqIdxPerToken == nullptr || (params.seqLens != nullptr && params.kvLens != nullptr), + "draft sample glue: the DSA metadata needs the sequence and KV lengths"); + tensorrt_llm::common::launchWithPdlWhenEnabled( + "draftSampleGlue", draftSampleGlueKernel, dim3(1), dim3(kBlockSize), 0, stream, params); +} + +void invokeDraftEhprojTable(DraftEhprojTableParams const& params, cudaStream_t stream) +{ + DraftSampleGlueParams const& s = params.step; + TLLM_CHECK_WITH_INFO(s.numRows > 0 && s.numRows <= kDraftSampleGlueMaxRows, + "draft eh_proj table serves 1 to %d rows, got %d", kDraftSampleGlueMaxRows, s.numRows); + TLLM_CHECK_WITH_INFO(draftEhprojTableSupportsTpSize(s.tpSize), + "draft eh_proj table serves 2 to %d ranks with whole 8-column shares of the hidden size, got %d", + kDraftEhprojTableMaxRanks, s.tpSize); + TLLM_CHECK_WITH_INFO(s.vocabShard > 0 && s.vocabShard % 8 == 0 && s.logitsStride % 8 == 0, + "draft eh_proj table needs 16-byte aligned rows of a multiple of 8 logits, got %d", s.vocabShard); + TLLM_CHECK_WITH_INFO(params.hiddenHalf != nullptr && params.hiddenHalfStride % 8 == 0 + && reinterpret_cast(params.hiddenHalf) % 32 == 0, + "draft eh_proj table needs 32-byte aligned rows of FP32 hidden-half outputs"); + TLLM_CHECK_WITH_INFO(s.glueOut == nullptr && s.nextNewTokens == nullptr, + "draft eh_proj table: no eh_proj input window and no next-new-tokens rows"); + TLLM_CHECK_WITH_INFO(params.tokenTable != nullptr && params.normWeight != nullptr && params.normedOut != nullptr + && params.residualOut != nullptr && params.outStride % 8 == 0, + "draft eh_proj table: missing token table, norm weight or 16-byte aligned output rows"); + TLLM_CHECK_WITH_INFO(s.reqIdxPerToken == nullptr || (s.seqLens != nullptr && s.kvLens != nullptr), + "draft eh_proj table: the DSA metadata needs the sequence and KV lengths"); + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kDraftEhprojTableClusterSize; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL() ? 1 : 0; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(kDraftEhprojTableClusterSize); + config.blockDim = dim3(kTableThreads); + config.dynamicSmemBytes = 0; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = 2; + // Up to TP8 the kernel keeps eight rank slots (registers of the polled rows); TP16 takes sixteen. + auto* const kernel = s.tpSize <= kTableTp8RankSlots ? &draftEhprojTableKernel + : &draftEhprojTableKernel; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, params)); +} + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.h new file mode 100644 index 000000000000..25920a5f26fb --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.h @@ -0,0 +1,188 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +//! Rows (requests) one launch serves. +constexpr int kDraftSampleGlueMaxRows = 8; +//! Hidden size of the eh_proj input glue (FlashInfer RMSNormKernel geometry: 64 threads x 8-element vectors x 12). +constexpr int kDraftSampleGlueHidden = 6144; +//! Bytes of one exchanged (index, value) packet. +constexpr int kDraftSampleGluePacketBytes = 16; + +//! How the next step's KV lengths follow from the current ones. +enum class DraftSampleGlueKvUpdate : int +{ + kNone = 0, //!< unchanged + kIncrement = 1, //!< kvLens[r] += 1 + kRewind = 2, //!< kvLens[r] -= draftLen - kvRewindAccepted[r] +}; + +//! One MTP draft step after the vocabulary-sharded draft LM head, in one CTA: +//! - the greedy draft token: the local argmax of this rank's logits shard, exchanged with the other ranks through a +//! dedicated MNNVL Lamport workspace (layout and flags of the MNNVL all-gather), and the global argmax of the pairs; +//! - the state of the next draft step: position ids, sequence and KV lengths, the DSA decode metadata of +//! fused_dsa_decode_metadata; +//! - the next step's eh_proj input: this rank's columns of cat([enorm(embed(token)), hnorm(hidden)], -1); +//! - after the last draft step, the next-new-tokens rows of the output. +//! Every part except the token is optional (skipped when its output pointer is null). +struct DraftSampleGlueParams +{ + // Draft logits of this rank, [numRows, vocabShard] BF16 with row stride logitsStride (elements). + __nv_bfloat16 const* logits{nullptr}; + int64_t logitsStride{0}; + int numRows{0}; + int vocabShard{0}; + int rank{0}; + int tpSize{1}; + + // Lamport workspace: three buffers of kDraftSampleGlueMaxRows * tpSize packets, the local unicast view, the + // multicast view and the nine 32-bit flags of the MNNVL all-gather. + void* exchangeLocal{nullptr}; + void* exchangeMcast{nullptr}; + uint32_t* exchangeFlags{nullptr}; + + // Draft tokens [numRows, numDraftTokens] int32; this step writes column draftStep. + int32_t* draftTokens{nullptr}; + int numDraftTokens{1}; + int draftStep{0}; + + // Next step's eh_proj input (skipped when glueOut is null): columns [colLo, colHi) of + // cat([enorm(embedTable[token]), hnorm(hidden)], -1), stored at glueOut[row * glueOutStride + col - colLo]. + __nv_bfloat16 const* embedTable{nullptr}; + int64_t vocabSize{0}; + __nv_bfloat16 const* hidden{nullptr}; + int64_t hiddenStride{0}; + __nv_bfloat16 const* enormWeight{nullptr}; + __nv_bfloat16 const* hnormWeight{nullptr}; + float eps{0.F}; + int colLo{0}; + int colHi{0}; + __nv_bfloat16* glueOut{nullptr}; + int64_t glueOutStride{0}; + + // Next step's position ids (skipped when positionsOut is null): + // positionsOut[r] = positionsIn[gatherIds ? gatherIds[r] : r] + 1, int32 or int64. + void const* positionsIn{nullptr}; + void* positionsOut{nullptr}; + bool positions64{false}; + int64_t const* gatherIds{nullptr}; + + // Sequence and KV lengths of the next step: seqLens[r] = 1 when fillSeqLens; kvLens[r] updated as kvUpdate says. + int32_t* seqLens{nullptr}; + bool fillSeqLens{false}; + int32_t* kvLens{nullptr}; + DraftSampleGlueKvUpdate kvUpdate{DraftSampleGlueKvUpdate::kNone}; + int32_t const* kvRewindAccepted{nullptr}; + int draftLen{0}; + + // DSA decode metadata of the next step (skipped when reqIdxPerToken is null), from the sequence and KV lengths + // above: what fused_dsa_decode_metadata writes for numRows requests and numTokens tokens. maxQueryLanes is its + // per-request token lanes (the power of two at or above the maximum query length). + int32_t const* blockOffsets{nullptr}; + int64_t blockOffsetsStride0{0}; + int64_t blockOffsetsStride1{1}; + int maxBlocks{0}; + int numTokens{0}; + int maxQueryLanes{1}; + int tokensPerBlock{0}; + int64_t dataBytesPerToken{0}; + int64_t scaleSize{0}; + int64_t blockStride{0}; + int64_t scaleBaseOffset{0}; + int32_t* reqIdxPerToken{nullptr}; + int64_t* slotMappingFp8{nullptr}; + int64_t* slotMappingScale{nullptr}; + int64_t* genKvIndptr{nullptr}; + int64_t* genCachedTokenIndptr{nullptr}; + + // Output rows after the last draft step (skipped when nextNewTokens is null), [numRows, numDraftTokens + 1] + // int32: nextNewTokens[r, 0] = acceptedTokens[batchIndices[r], numAccepted[r] - 1], then the draft tokens. + int32_t const* acceptedTokens{nullptr}; + int64_t acceptedTokensStride{0}; + int32_t const* batchIndices{nullptr}; + int32_t const* numAccepted{nullptr}; + int32_t* nextNewTokens{nullptr}; +}; + +void invokeDraftSampleGlue(DraftSampleGlueParams const& params, cudaStream_t stream); + +//! CTAs of the cluster that runs invokeDraftEhprojTable; each owns kDraftSampleGlueHidden / 8 hidden columns. +constexpr int kDraftEhprojTableClusterSize = 8; +//! Largest TP size of invokeDraftEhprojTable (2, 4, 8 or 16 ranks: each rank's share of the hidden-half outputs is +//! whole 8-column vectors). +constexpr int kDraftEhprojTableMaxRanks = 16; + +//! Whether invokeDraftEhprojTable serves tpSize ranks. +constexpr bool draftEhprojTableSupportsTpSize(int tpSize) +{ + return tpSize >= 2 && tpSize <= kDraftEhprojTableMaxRanks && kDraftSampleGlueHidden % (8 * tpSize) == 0; +} + +//! Bytes one rank publishes per row in invokeDraftEhprojTable: its (index, value) packet, the BF16 token-table row of +//! its candidate token and its kDraftSampleGlueHidden / tpSize FP32 hidden-half eh_proj outputs. +constexpr int draftEhprojTableSlotBytes(int tpSize) +{ + return kDraftSampleGluePacketBytes + 2 * kDraftSampleGlueHidden + 4 * kDraftSampleGlueHidden / tpSize; +} + +//! One MTP draft step after the vocabulary-sharded draft LM head that also replaces the next step's eh_proj and its +//! all-reduce + input_layernorm, in one cluster of kDraftEhprojTableClusterSize CTAs: +//! - the greedy draft token and the next step's state, as invokeDraftSampleGlue computes them; +//! - the next step's eh_proj output x = BF16(tokenTable[token] + hiddenHalf), one FP32 add, where tokenTable holds +//! eh_proj's embedding half applied to enorm(embed(t)) for the tokens t of this rank's vocabulary shard (BF16, one +//! rounding of the FP32 sum), and hiddenHalf is eh_proj's hidden half applied to hnorm(hidden) in FP32, the ranks +//! computing H / tpSize outputs each. Every rank publishes the token-table row of its candidate and its share of +//! hiddenHalf with its (index, value) packet in one exchange; every rank then takes the winner's row; +//! - input_layernorm(x) and x, what eh_proj's all-reduce with input_layernorm and a zero residual returns. +struct DraftEhprojTableParams +{ + //! Token, state and exchange workspace as for invokeDraftSampleGlue; glueOut and nextNewTokens must be null. The + //! workspace buffers hold numRows * tpSize slots of draftEhprojTableSlotBytes(tpSize) bytes. + DraftSampleGlueParams step; + //! This rank's hiddenHalf columns [rank * H / tpSize, (rank + 1) * H / tpSize), FP32 [numRows, H / tpSize] with + //! row stride hiddenHalfStride. + float const* hiddenHalf{nullptr}; + int64_t hiddenHalfStride{0}; + //! [vocabShard, H]: row t is eh_proj's embedding half applied to enorm(embed(rank * vocabShard + t)). + __nv_bfloat16 const* tokenTable{nullptr}; + //! input_layernorm weight [H] and epsilon. + __nv_bfloat16 const* normWeight{nullptr}; + float normEps{0.F}; + //! Outputs [numRows, H] with row stride outStride: input_layernorm(x) and x. + __nv_bfloat16* normedOut{nullptr}; + __nv_bfloat16* residualOut{nullptr}; + int64_t outStride{0}; +}; + +void invokeDraftEhprojTable(DraftEhprojTableParams const& params, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 8a84a5bf4724..a2d6c763597b 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -145,7 +145,8 @@ add_th_common_family( userbuffersFinalizeOp.cpp userbuffersTensor.cpp virtualMemoryAllocator.cpp - fusedGemmAllreduceOp.cpp) + fusedGemmAllreduceOp.cpp + glmDraftSampleGlueOp.cpp) add_th_common_family(th_common_sequence causalConv1dOp.cpp mambaConv1dOp.cpp mamba2MTPSSMCacheOp.cpp selectiveScanOp.cpp) diff --git a/cpp/tensorrt_llm/thop/glmDraftSampleGlueOp.cpp b/cpp/tensorrt_llm/thop/glmDraftSampleGlueOp.cpp new file mode 100644 index 000000000000..27390af30f41 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmDraftSampleGlueOp.cpp @@ -0,0 +1,366 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/mcastDevMemUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.h" +#include "tensorrt_llm/runtime/mcastDeviceMemory.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace tk = tensorrt_llm::kernels::glm_moe_dsa; + +bool isAligned16(at::Tensor const& t) +{ + return reinterpret_cast(t.data_ptr()) % 16 == 0; +} + +// The logits, exchange workspace and draft-token fields of both draft-step ops; slotBytes is what one rank publishes +// per row. Returns the number of rows. +int64_t setDraftStepExchange(tk::DraftSampleGlueParams& params, char const* op, at::Tensor const& logits, + at::Tensor& exchange_buffer, at::Tensor& exchange_flags, int64_t group_size, at::Tensor& draft_tokens, + int64_t draft_step, int64_t slotBytes) +{ + auto* mcastMem = tensorrt_llm::common::findMcastDevMemBuffer(exchange_buffer.data_ptr()); + TORCH_CHECK(mcastMem != nullptr, op, ": exchange_buffer must come from a McastGPUBuffer"); + TORCH_CHECK(static_cast(mcastMem->getWorldSize()) == group_size, op, ": group_size ", group_size, + " does not match the workspace group of ", mcastMem->getWorldSize(), " ranks"); + TORCH_CHECK(exchange_flags.element_size() == 4 && exchange_flags.numel() >= 9, op, + ": exchange_flags must hold the nine 32-bit Lamport flags"); + CHECK_INPUT(logits, at::ScalarType::BFloat16); + TORCH_CHECK(logits.dim() == 2 && logits.size(0) >= 1 && logits.size(0) <= tk::kDraftSampleGlueMaxRows, op, + ": logits must be [rows, vocab_shard] with 1 to ", tk::kDraftSampleGlueMaxRows, " rows"); + int64_t const rows = logits.size(0); + TORCH_CHECK(logits.size(1) % 8 == 0 && isAligned16(logits), op, + ": logits rows must be 16-byte aligned and a multiple of 8 columns"); + int64_t const bufferBytes = exchange_buffer.numel() * exchange_buffer.element_size() / 3; + TORCH_CHECK(bufferBytes >= rows * group_size * slotBytes, op, ": the Lamport buffer is too small for ", rows, + " rows"); + CHECK_INPUT(draft_tokens, at::ScalarType::Int); + TORCH_CHECK(draft_tokens.dim() == 2 && draft_tokens.size(0) == rows && draft_step >= 0 + && draft_step < draft_tokens.size(1), + op, ": draft_tokens must be [rows, draft_len] with draft_step inside"); + + params.logits = reinterpret_cast<__nv_bfloat16 const*>(logits.data_ptr()); + params.logitsStride = logits.stride(0); + params.numRows = static_cast(rows); + params.vocabShard = static_cast(logits.size(1)); + params.rank = static_cast(mcastMem->getRank()); + params.tpSize = static_cast(group_size); + params.exchangeLocal = exchange_buffer.data_ptr(); + params.exchangeMcast = mcastMem->getMulticastPtr(); + params.exchangeFlags = reinterpret_cast(exchange_flags.data_ptr()); + params.draftTokens = draft_tokens.data_ptr(); + params.numDraftTokens = static_cast(draft_tokens.size(1)); + params.draftStep = static_cast(draft_step); + return rows; +} + +// The next step's position ids, sequence / KV lengths and DSA decode metadata of both draft-step ops. +void setDraftStepState(tk::DraftSampleGlueParams& params, char const* op, int64_t rows, + std::optional const& positions_in, std::optional& positions_out, + std::optional const& gather_ids, std::optional& seq_lens, bool fill_seq_lens, + std::optional& kv_lens, int64_t kv_update, std::optional const& rewind_accepted, + int64_t draft_len, std::optional const& block_offsets, std::vector const& dsa_layout, + std::optional& req_idx_per_token, std::optional& slot_mapping_fp8, + std::optional& slot_mapping_scale, std::optional& gen_kv_indptr, + std::optional& gen_cached_token_indptr) +{ + if (positions_out.has_value()) + { + TORCH_CHECK(positions_in.has_value(), op, ": positions_out needs positions_in"); + auto const posType = positions_in->scalar_type(); + TORCH_CHECK(posType == at::ScalarType::Int || posType == at::ScalarType::Long, op, + ": position ids must be int32 or int64"); + CHECK_INPUT(positions_in.value(), posType); + CHECK_INPUT(positions_out.value(), posType); + TORCH_CHECK(positions_out->numel() >= rows, op, ": positions_out holds one id per row"); + params.positionsIn = positions_in->data_ptr(); + params.positionsOut = positions_out->data_ptr(); + params.positions64 = posType == at::ScalarType::Long; + if (gather_ids.has_value()) + { + CHECK_INPUT(gather_ids.value(), at::ScalarType::Long); + TORCH_CHECK(gather_ids->numel() >= rows, op, ": gather_ids holds one index per row"); + params.gatherIds = gather_ids->data_ptr(); + } + } + + if (seq_lens.has_value()) + { + CHECK_INPUT(seq_lens.value(), at::ScalarType::Int); + TORCH_CHECK(seq_lens->numel() >= rows, op, ": seq_lens holds one length per row"); + params.seqLens = seq_lens->data_ptr(); + params.fillSeqLens = fill_seq_lens; + } + if (kv_lens.has_value()) + { + CHECK_INPUT(kv_lens.value(), at::ScalarType::Int); + TORCH_CHECK(kv_lens->numel() >= rows, op, ": kv_lens holds one length per row"); + TORCH_CHECK(kv_update >= 0 && kv_update <= 2, op, ": unknown kv_update ", kv_update); + params.kvLens = kv_lens->data_ptr(); + params.kvUpdate = static_cast(kv_update); + if (params.kvUpdate == tk::DraftSampleGlueKvUpdate::kRewind) + { + TORCH_CHECK(rewind_accepted.has_value(), op, ": kv_update 2 needs rewind_accepted"); + CHECK_INPUT(rewind_accepted.value(), at::ScalarType::Int); + TORCH_CHECK(rewind_accepted->numel() >= rows, op, ": rewind_accepted holds one count per row"); + params.kvRewindAccepted = rewind_accepted->data_ptr(); + } + params.draftLen = static_cast(draft_len); + } + + if (req_idx_per_token.has_value()) + { + TORCH_CHECK(seq_lens.has_value() && kv_lens.has_value() && block_offsets.has_value() + && slot_mapping_fp8.has_value() && slot_mapping_scale.has_value() && gen_kv_indptr.has_value() + && gen_cached_token_indptr.has_value() && dsa_layout.size() == 7, + op, ": the DSA metadata needs the lengths, the block offsets, all five outputs and the layout"); + CHECK_TH_CUDA(block_offsets.value()); + TORCH_CHECK(block_offsets->scalar_type() == at::ScalarType::Int && block_offsets->dim() == 2 + && block_offsets->size(0) >= rows, + op, ": block_offsets must be int32 [rows, max_blocks]"); + int64_t const numTokens = dsa_layout[0]; + CHECK_INPUT(req_idx_per_token.value(), at::ScalarType::Int); + CHECK_INPUT(slot_mapping_fp8.value(), at::ScalarType::Long); + CHECK_INPUT(slot_mapping_scale.value(), at::ScalarType::Long); + CHECK_INPUT(gen_kv_indptr.value(), at::ScalarType::Long); + CHECK_INPUT(gen_cached_token_indptr.value(), at::ScalarType::Long); + TORCH_CHECK(req_idx_per_token->numel() >= numTokens && slot_mapping_fp8->numel() >= numTokens + && slot_mapping_scale->numel() >= numTokens && gen_kv_indptr->numel() >= rows + 1 + && gen_cached_token_indptr->numel() >= rows + 1, + op, ": DSA metadata outputs too small"); + params.blockOffsets = block_offsets->data_ptr(); + params.blockOffsetsStride0 = block_offsets->stride(0); + params.blockOffsetsStride1 = block_offsets->stride(1); + params.maxBlocks = static_cast(block_offsets->size(1)); + params.numTokens = static_cast(numTokens); + params.maxQueryLanes = static_cast(dsa_layout[1]); + params.tokensPerBlock = static_cast(dsa_layout[2]); + params.dataBytesPerToken = dsa_layout[3]; + params.scaleSize = dsa_layout[4]; + params.blockStride = dsa_layout[5]; + params.scaleBaseOffset = dsa_layout[6]; + params.reqIdxPerToken = req_idx_per_token->data_ptr(); + params.slotMappingFp8 = slot_mapping_fp8->data_ptr(); + params.slotMappingScale = slot_mapping_scale->data_ptr(); + params.genKvIndptr = gen_kv_indptr->data_ptr(); + params.genCachedTokenIndptr = gen_cached_token_indptr->data_ptr(); + } +} + +} // namespace + +// One MTP draft step after the vocabulary-sharded draft LM head (see draftSampleGlue.h): the greedy draft token of +// every row into column draft_step of draft_tokens, through the dedicated MNNVL Lamport workspace exchange_buffer / +// exchange_flags; optionally the next step's eh_proj input window (glue_out), position ids, sequence and KV lengths +// (kv_update 0: none, 1: + 1, 2: - (draft_len - rewind_accepted)), the DSA decode metadata (dsa_layout = [num_tokens, +// max_query_lanes, tokens_per_block, data_bytes_per_token, scale_size, block_stride, scale_base_offset]) and the +// next-new-tokens rows after the last step. +void glmDraftSampleGlue(at::Tensor const& logits, at::Tensor& exchange_buffer, at::Tensor& exchange_flags, + int64_t group_size, at::Tensor& draft_tokens, int64_t draft_step, std::optional const& embed_table, + std::optional const& hidden, std::optional const& enorm_weight, + std::optional const& hnorm_weight, double eps, int64_t col_lo, int64_t col_hi, + std::optional glue_out, std::optional const& positions_in, + std::optional positions_out, std::optional const& gather_ids, + std::optional seq_lens, bool fill_seq_lens, std::optional kv_lens, int64_t kv_update, + std::optional const& rewind_accepted, int64_t draft_len, std::optional const& block_offsets, + std::vector const& dsa_layout, std::optional req_idx_per_token, + std::optional slot_mapping_fp8, std::optional slot_mapping_scale, + std::optional gen_kv_indptr, std::optional gen_cached_token_indptr, + std::optional const& accepted_tokens, std::optional const& batch_indices, + std::optional const& num_accepted, std::optional next_new_tokens) +{ + tk::DraftSampleGlueParams params; + int64_t const rows = setDraftStepExchange(params, "glm_draft_sample_glue", logits, exchange_buffer, exchange_flags, + group_size, draft_tokens, draft_step, tk::kDraftSampleGluePacketBytes); + + if (glue_out.has_value()) + { + constexpr int64_t kHidden = tk::kDraftSampleGlueHidden; + TORCH_CHECK(col_lo >= 0 && col_lo < col_hi && col_hi <= 2 * kHidden && col_lo % 8 == 0 && col_hi % 8 == 0, + "glm_draft_sample_glue: invalid eh_proj input window [", col_lo, ", ", col_hi, ")"); + TORCH_CHECK( + enorm_weight.has_value() && hnorm_weight.has_value(), "glm_draft_sample_glue: norm weights missing"); + for (at::Tensor const& weight : {enorm_weight.value(), hnorm_weight.value()}) + { + CHECK_INPUT(weight, at::ScalarType::BFloat16); + TORCH_CHECK(weight.numel() == kHidden && isAligned16(weight), + "glm_draft_sample_glue: norm weights must be [6144]"); + } + TORCH_CHECK(glue_out->scalar_type() == at::ScalarType::BFloat16 && glue_out->dim() == 2 + && glue_out->size(0) == rows && glue_out->size(1) == col_hi - col_lo && glue_out->stride(1) == 1 + && glue_out->stride(0) % 8 == 0 && isAligned16(*glue_out), + "glm_draft_sample_glue: glue_out must be a 16-byte aligned BF16 [rows, col_hi - col_lo] view"); + if (col_lo < kHidden) + { + TORCH_CHECK(embed_table.has_value(), "glm_draft_sample_glue: the window needs the embedding table"); + CHECK_INPUT(embed_table.value(), at::ScalarType::BFloat16); + TORCH_CHECK(embed_table->dim() == 2 && embed_table->size(1) == kHidden && isAligned16(*embed_table), + "glm_draft_sample_glue: embed_table must be [vocab, 6144]"); + params.embedTable = reinterpret_cast<__nv_bfloat16 const*>(embed_table->data_ptr()); + params.vocabSize = embed_table->size(0); + } + if (col_hi > kHidden) + { + TORCH_CHECK(hidden.has_value(), "glm_draft_sample_glue: the window needs the hidden rows"); + TORCH_CHECK(hidden->scalar_type() == at::ScalarType::BFloat16 && hidden->dim() == 2 + && hidden->size(0) == rows && hidden->size(1) == kHidden && hidden->stride(1) == 1 + && hidden->stride(0) % 8 == 0 && isAligned16(*hidden), + "glm_draft_sample_glue: hidden must be a 16-byte aligned BF16 [rows, 6144] view"); + params.hidden = reinterpret_cast<__nv_bfloat16 const*>(hidden->data_ptr()); + params.hiddenStride = hidden->stride(0); + } + params.enormWeight = reinterpret_cast<__nv_bfloat16 const*>(enorm_weight->data_ptr()); + params.hnormWeight = reinterpret_cast<__nv_bfloat16 const*>(hnorm_weight->data_ptr()); + params.eps = static_cast(eps); + params.colLo = static_cast(col_lo); + params.colHi = static_cast(col_hi); + params.glueOut = reinterpret_cast<__nv_bfloat16*>(glue_out->data_ptr()); + params.glueOutStride = glue_out->stride(0); + } + + setDraftStepState(params, "glm_draft_sample_glue", rows, positions_in, positions_out, gather_ids, seq_lens, + fill_seq_lens, kv_lens, kv_update, rewind_accepted, draft_len, block_offsets, dsa_layout, req_idx_per_token, + slot_mapping_fp8, slot_mapping_scale, gen_kv_indptr, gen_cached_token_indptr); + + if (next_new_tokens.has_value()) + { + TORCH_CHECK(accepted_tokens.has_value() && batch_indices.has_value() && num_accepted.has_value(), + "glm_draft_sample_glue: next_new_tokens needs the accepted tokens, batch indices and counts"); + CHECK_INPUT(next_new_tokens.value(), at::ScalarType::Int); + CHECK_TH_CUDA(accepted_tokens.value()); + TORCH_CHECK(accepted_tokens->scalar_type() == at::ScalarType::Int && accepted_tokens->dim() == 2 + && accepted_tokens->stride(1) == 1, + "glm_draft_sample_glue: accepted_tokens must be int32 [batch, draft_len + 1] rows"); + CHECK_INPUT(batch_indices.value(), at::ScalarType::Int); + CHECK_INPUT(num_accepted.value(), at::ScalarType::Int); + TORCH_CHECK(next_new_tokens->numel() >= rows * (draft_tokens.size(1) + 1) && batch_indices->numel() >= rows + && num_accepted->numel() >= rows, + "glm_draft_sample_glue: next_new_tokens must be [rows, draft_len + 1]"); + params.acceptedTokens = accepted_tokens->data_ptr(); + params.acceptedTokensStride = accepted_tokens->stride(0); + params.batchIndices = batch_indices->data_ptr(); + params.numAccepted = num_accepted->data_ptr(); + params.nextNewTokens = next_new_tokens->data_ptr(); + } + + tk::invokeDraftSampleGlue(params, at::cuda::getCurrentCUDAStream(logits.get_device())); +} + +// One MTP draft step that also replaces the next step's eh_proj with its all-reduce and input_layernorm (see +// invokeDraftEhprojTable): the draft token and the next step's state as glm_draft_sample_glue computes them, +// residual_out = x = BF16(token_table[token] + hidden_half), one FP32 add, and normed_out = input_layernorm(x) with +// norm_weight and eps. hidden_half is this rank's FP32 [rows, 6144 / group_size] share of eh_proj's hidden half, +// token_table this rank's [vocab_shard, 6144] rows; exchange_buffer / exchange_flags are a dedicated MNNVL Lamport +// workspace of draftEhprojTableSlotBytes(group_size) bytes per rank and row. +void glmDraftEhprojTable(at::Tensor const& logits, at::Tensor& exchange_buffer, at::Tensor& exchange_flags, + int64_t group_size, at::Tensor& draft_tokens, int64_t draft_step, at::Tensor const& hidden_half, + at::Tensor const& token_table, at::Tensor const& norm_weight, double eps, at::Tensor& normed_out, + at::Tensor& residual_out, std::optional const& positions_in, std::optional positions_out, + std::optional const& gather_ids, std::optional seq_lens, bool fill_seq_lens, + std::optional kv_lens, int64_t kv_update, std::optional const& rewind_accepted, + int64_t draft_len, std::optional const& block_offsets, std::vector const& dsa_layout, + std::optional req_idx_per_token, std::optional slot_mapping_fp8, + std::optional slot_mapping_scale, std::optional gen_kv_indptr, + std::optional gen_cached_token_indptr) +{ + constexpr int64_t kHidden = tk::kDraftSampleGlueHidden; + TORCH_CHECK(tk::draftEhprojTableSupportsTpSize(static_cast(group_size)), + "glm_draft_ehproj_table: group_size must be 2 to ", tk::kDraftEhprojTableMaxRanks, + " ranks with whole 8-column shares of the hidden size, got ", group_size); + tk::DraftEhprojTableParams params; + int64_t const rows = setDraftStepExchange(params.step, "glm_draft_ehproj_table", logits, exchange_buffer, + exchange_flags, group_size, draft_tokens, draft_step, tk::draftEhprojTableSlotBytes(group_size)); + + TORCH_CHECK(hidden_half.is_cuda() && hidden_half.scalar_type() == at::ScalarType::Float && hidden_half.dim() == 2 + && hidden_half.size(0) == rows && hidden_half.size(1) == kHidden / group_size && hidden_half.stride(1) == 1 + && hidden_half.stride(0) % 8 == 0 && reinterpret_cast(hidden_half.data_ptr()) % 32 == 0, + "glm_draft_ehproj_table: hidden_half must be a 32-byte aligned FP32 [rows, 6144 / group_size] view"); + CHECK_INPUT(token_table, at::ScalarType::BFloat16); + TORCH_CHECK(token_table.dim() == 2 && token_table.size(0) == logits.size(1) && token_table.size(1) == kHidden + && isAligned16(token_table), + "glm_draft_ehproj_table: token_table must be [vocab_shard, 6144]"); + CHECK_INPUT(norm_weight, at::ScalarType::BFloat16); + TORCH_CHECK(norm_weight.numel() == kHidden && isAligned16(norm_weight), + "glm_draft_ehproj_table: norm_weight must be [6144]"); + for (at::Tensor const& out : {normed_out, residual_out}) + { + TORCH_CHECK(out.is_cuda() && out.scalar_type() == at::ScalarType::BFloat16 && out.dim() == 2 + && out.size(0) == rows && out.size(1) == kHidden && out.stride(1) == 1 + && out.stride(0) == normed_out.stride(0) && out.stride(0) % 8 == 0 && isAligned16(out), + "glm_draft_ehproj_table: normed_out and residual_out must be 16-byte aligned BF16 [rows, 6144] views with " + "one row stride"); + } + params.hiddenHalf = hidden_half.data_ptr(); + params.hiddenHalfStride = hidden_half.stride(0); + params.tokenTable = reinterpret_cast<__nv_bfloat16 const*>(token_table.data_ptr()); + params.normWeight = reinterpret_cast<__nv_bfloat16 const*>(norm_weight.data_ptr()); + params.normEps = static_cast(eps); + params.normedOut = reinterpret_cast<__nv_bfloat16*>(normed_out.data_ptr()); + params.residualOut = reinterpret_cast<__nv_bfloat16*>(residual_out.data_ptr()); + params.outStride = normed_out.stride(0); + + setDraftStepState(params.step, "glm_draft_ehproj_table", rows, positions_in, positions_out, gather_ids, seq_lens, + fill_seq_lens, kv_lens, kv_update, rewind_accepted, draft_len, block_offsets, dsa_layout, req_idx_per_token, + slot_mapping_fp8, slot_mapping_scale, gen_kv_indptr, gen_cached_token_indptr); + + tk::invokeDraftEhprojTable(params, at::cuda::getCurrentCUDAStream(logits.get_device())); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_draft_sample_glue(Tensor logits, Tensor(a!) exchange_buffer, Tensor(b!) exchange_flags, int group_size, " + "Tensor(c!) draft_tokens, int draft_step, Tensor? embed_table, Tensor? hidden, Tensor? enorm_weight, " + "Tensor? hnorm_weight, float eps, int col_lo, int col_hi, Tensor(d!)? glue_out, Tensor? positions_in, " + "Tensor(e!)? positions_out, Tensor? gather_ids, Tensor(f!)? seq_lens, bool fill_seq_lens, " + "Tensor(g!)? kv_lens, int kv_update, Tensor? rewind_accepted, int draft_len, Tensor? block_offsets, " + "int[] dsa_layout, Tensor(h!)? req_idx_per_token, Tensor(i!)? slot_mapping_fp8, " + "Tensor(j!)? slot_mapping_scale, Tensor(k!)? gen_kv_indptr, Tensor(l!)? gen_cached_token_indptr, " + "Tensor? accepted_tokens, Tensor? batch_indices, Tensor? num_accepted, Tensor(m!)? next_new_tokens) -> ()"); + m.def( + "glm_draft_ehproj_table(Tensor logits, Tensor(a!) exchange_buffer, Tensor(b!) exchange_flags, int group_size, " + "Tensor(c!) draft_tokens, int draft_step, Tensor hidden_half, Tensor token_table, Tensor norm_weight, " + "float eps, Tensor(d!) normed_out, Tensor(e!) residual_out, Tensor? positions_in, Tensor(f!)? positions_out, " + "Tensor? gather_ids, Tensor(g!)? seq_lens, bool fill_seq_lens, Tensor(h!)? kv_lens, int kv_update, " + "Tensor? rewind_accepted, int draft_len, Tensor? block_offsets, int[] dsa_layout, " + "Tensor(i!)? req_idx_per_token, Tensor(j!)? slot_mapping_fp8, Tensor(k!)? slot_mapping_scale, " + "Tensor(l!)? gen_kv_indptr, Tensor(m!)? gen_cached_token_indptr) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_draft_sample_glue", &tensorrt_llm::torch_ext::glmDraftSampleGlue); + m.impl("glm_draft_ehproj_table", &tensorrt_llm::torch_ext::glmDraftEhprojTable); +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_exchange.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_exchange.py new file mode 100644 index 000000000000..cdc715a2f082 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_exchange.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MNNVL Lamport workspace the two draft-step ops exchange through, and their shape constants. + +``glm_draft_sample_glue`` and ``glm_draft_ehproj_table`` publish each rank's (index, value) packet (and, for the +table op, a token-table row and the rank's hidden-half outputs) through the multicast address of a dedicated +workspace and poll the local copy until every rank's slot has landed. The workspace has three Lamport buffers and +the nine 32-bit flags of the MNNVL all-reduce layout (current index, dirty index, bytes per buffer, dirty stages, +four byte counts to clear, an access counter); an unwritten word holds 0xFFFFFFFF rather than the all-reduce's -0.0. + +Each op empties the buffer of its previous call with the threads of its current launch, so a workspace serves one +op and nothing else. +""" + +from __future__ import annotations + +import os +from typing import Any + +import torch + +from tensorrt_llm._torch.distributed.ops import _get_mnnvl_workspace_comm, _mnnvl_device_index +from tensorrt_llm._utils import torch_pybind11_abi +from tensorrt_llm.bindings.internal.runtime import McastGPUBuffer +from tensorrt_llm.mapping import Mapping + +#: Rows (requests) one draft-step launch serves (kDraftSampleGlueMaxRows). +MAX_ROWS = 8 +#: Bytes of one exchanged (index, value) packet (kDraftSampleGluePacketBytes). +PACKET_BYTES = 16 +#: Hidden size the draft-step kernels are built for (kDraftSampleGlueHidden). +HIDDEN_SIZE = 6144 +#: Largest TP size of glm_draft_ehproj_table (kDraftEhprojTableMaxRanks). +EHPROJ_TABLE_MAX_RANKS = 16 + + +def ehproj_table_supports_tp(tp_size: int) -> bool: + """Whether glm_draft_ehproj_table serves ``tp_size`` ranks (draftEhprojTableSupportsTpSize). + + 2 to 16 ranks whose shares of the hidden size are whole 8-column vectors. + """ + return 2 <= tp_size <= EHPROJ_TABLE_MAX_RANKS and HIDDEN_SIZE % (8 * tp_size) == 0 + + +def sample_glue_slot_bytes(tp_size: int) -> int: + """Bytes one rank publishes per row in glm_draft_sample_glue: its (index, value) packet.""" + del tp_size + return PACKET_BYTES + + +def ehproj_table_slot_bytes(tp_size: int) -> int: + """Bytes one rank publishes per row in glm_draft_ehproj_table (draftEhprojTableSlotBytes). + + The packet, the BF16 token-table row of its candidate and its FP32 share of the hidden-half outputs. + """ + return PACKET_BYTES + 2 * HIDDEN_SIZE + 4 * HIDDEN_SIZE // tp_size + + +def _all_succeeded(comm: Any, ok: bool) -> bool: + """Whether ``ok`` holds on every rank of ``comm`` (an MPI communicator or a ProcessGroup with a CPU backend).""" + if isinstance(comm, torch.distributed.ProcessGroup): + flag = torch.tensor([1 if ok else 0], dtype=torch.int32) + work = comm._get_backend(torch.device("cpu")).allreduce([flag]) + work.wait() + return int(flag.item()) == comm.size() + return comm.allreduce(int(ok)) == comm.Get_size() + + +class DraftExchangeWorkspace: + """One draft-step op's MNNVL workspace over a TP group. + + ``buffer`` is the local unicast view of the three Lamport buffers (int32, every word 0xFFFFFFFF) and ``flags`` + the nine uint32 Lamport flags. Construction is collective over the group and either succeeds on every rank or + raises on every rank. Build it outside CUDA-graph capture, once per op; the captured graphs keep its addresses. + """ + + NUM_LAMPORT_BUFFERS = 3 + + def __init__( + self, + buffer_size_bytes: int, + group_size: int, + group_rank: int, + device_index: int, + use_fabric_handle: bool, + comm: Any, + ): + """Allocate ``3 * buffer_size_bytes`` of multicast memory over the group and reset it. + + ``comm`` is the group's MPI communicator or a ``torch.distributed.ProcessGroup`` with a CPU backend, ordered + by group rank. + """ + assert buffer_size_bytes > 0 and buffer_size_bytes % 16 == 0, buffer_size_bytes + self.buffer_size_bytes = buffer_size_bytes + self.group_size = group_size + workspace_size_bytes = self.NUM_LAMPORT_BUFFERS * buffer_size_bytes + error = None + state = None + try: + if isinstance(comm, torch.distributed.ProcessGroup): + handle = McastGPUBuffer( + workspace_size_bytes, + group_size, + group_rank, + device_index, + use_fabric_handle, + process_group=comm, + pybind11_abi=torch_pybind11_abi(), + ) + else: + handle = McastGPUBuffer( + workspace_size_bytes, + group_size, + group_rank, + device_index, + use_fabric_handle, + comm.py2f(), + ) + buffer = handle.get_uc_buffer(group_rank, (workspace_size_bytes // 4,), torch.int32, 0) + with torch.inference_mode(): + buffer.fill_(-1) + flags = torch.tensor( + [0, 2, buffer_size_bytes, 0, 0, 0, 0, 0, 0], + dtype=torch.uint32, + device=torch.device("cuda", device_index), + ) + torch.cuda.synchronize() + state = (handle, buffer, flags) + except (RuntimeError, ValueError) as e: + error = e + if not _all_succeeded(comm, error is None): + raise RuntimeError( + "draft-step MNNVL workspace construction failed on at least one rank" + ) from error + # The handle keeps the multicast allocation alive; the communicator outlives it. + self._handle, self.buffer, self.flags = state + self._comm = comm + + @classmethod + def for_mapping(cls, mapping: Mapping, buffer_size_bytes: int) -> "DraftExchangeWorkspace": + """The workspace over ``mapping``'s TP group, with the MNNVL all-reduce workspace's communicator. + + Fabric handles across nodes or under ``TRTLLM_FORCE_MNNVL_AR``, NVLS handles within a node, as the MNNVL + all-reduce workspace takes them. + """ + use_fabric_handle = ( + os.environ.get("TRTLLM_FORCE_MNNVL_AR", "0") == "1" or mapping.is_multi_node() + ) + return cls( + buffer_size_bytes, + mapping.tp_size, + mapping.tp_rank, + _mnnvl_device_index(mapping), + use_fabric_handle, + _get_mnnvl_workspace_comm(mapping), + ) + + @classmethod + def for_sample_glue(cls, mapping: Mapping) -> "DraftExchangeWorkspace": + """The workspace of glm_draft_sample_glue for up to MAX_ROWS rows.""" + slot_bytes = sample_glue_slot_bytes(mapping.tp_size) + return cls.for_mapping(mapping, MAX_ROWS * mapping.tp_size * slot_bytes) + + @classmethod + def for_ehproj_table(cls, mapping: Mapping) -> "DraftExchangeWorkspace": + """The workspace of glm_draft_ehproj_table for up to MAX_ROWS rows.""" + slot_bytes = ehproj_table_slot_bytes(mapping.tp_size) + return cls.for_mapping(mapping, MAX_ROWS * mapping.tp_size * slot_bytes) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_step_state.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_step_state.py new file mode 100644 index 000000000000..65ba0b8f088d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_step_state.py @@ -0,0 +1,189 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Plain-torch mirrors of what both draft-step ops write besides their main output. + +The greedy draft token over the vocabulary-sharded logits, and the next draft step's state: position ids, +sequence and KV lengths, and the DSA decode metadata that the fused DSA decode-metadata kernel writes for a +generation-only batch. Shared by the references of ``glm_draft_sample_glue`` and ``glm_draft_ehproj_table``. +""" + +from __future__ import annotations + +from typing import Dict, List, Optional, Sequence + +import torch + + +def global_argmax(per_rank_logits: Sequence[torch.Tensor]) -> torch.Tensor: + """The argmax of the rank-ordered concatenation of the logits shards, in torch.argmax's order. + + A NaN beats every number and the first NaN wins; otherwise the largest value, the lowest index on ties. + Returns int32 ``[rows]``. + """ + logits = torch.cat([shard.float() for shard in per_rank_logits], dim=-1) + return torch.argmax(logits, dim=-1).to(torch.int32) + + +def _rms_normalize(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """``bf16((x * rsqrt(mean(x^2) + eps)) * weight)`` in FP32, one row per leading index.""" + rstd = torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + eps) + return ((x.float() * rstd) * weight.float()).to(torch.bfloat16) + + +def eh_proj_input_window( + tokens: torch.Tensor, + embed_table: torch.Tensor, + hidden: torch.Tensor, + enorm_weight: torch.Tensor, + hnorm_weight: torch.Tensor, + eps: float, + col_lo: int, + col_hi: int, +) -> torch.Tensor: + """Columns ``[col_lo, col_hi)`` of ``cat([enorm(embed(tokens)), hnorm(hidden)], -1)``. + + An id outside the table gives a zero embedding row. + """ + vocab = embed_table.shape[0] + ids = tokens.long() + valid = (ids >= 0) & (ids < vocab) + embeds = embed_table[ids.clamp(0, vocab - 1)].float() * valid.unsqueeze(1) + concat = torch.cat( + [_rms_normalize(embeds, enorm_weight, eps), _rms_normalize(hidden, hnorm_weight, eps)], + dim=-1, + ) + return concat[:, col_lo:col_hi].contiguous() + + +def next_step_state( + rows: int, + positions_in: Optional[torch.Tensor], + positions_out: Optional[torch.Tensor], + gather_ids: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + fill_seq_lens: bool, + kv_lens: Optional[torch.Tensor], + kv_update: int, + rewind_accepted: Optional[torch.Tensor], + draft_len: int, + block_offsets: Optional[torch.Tensor], + dsa_layout: List[int], + req_idx_per_token: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + gen_kv_indptr: Optional[torch.Tensor], + gen_cached_token_indptr: Optional[torch.Tensor], +) -> Dict[str, torch.Tensor]: + """What a draft-step op leaves in each state buffer it is handed, read as the buffers were before the call. + + Position ids: ``positions_out[r] = positions_in[gather_ids[r] if gather_ids else r] + 1``. Sequence lengths: + filled with 1 when ``fill_seq_lens``. KV lengths: ``+ 1`` (kv_update 1) or ``- (draft_len - + rewind_accepted[r])`` (kv_update 2). DSA decode metadata, from those lengths: the generation KV and cached + indptrs and, for each request's tokens (at most ``dsa_layout[1]`` per request, below ``dsa_layout[0]`` in + all), the request index and the indexer-K cache slot offsets, the block index truncated and clamped as the + fused DSA decode-metadata kernel computes it. Entries the op does not write keep their values. Returns the + expected buffers by argument name. + """ + out: Dict[str, torch.Tensor] = {} + if positions_out is not None: + expected = positions_out.clone() + source = gather_ids[:rows].long() if gather_ids is not None else torch.arange(rows) + expected[:rows] = positions_in.reshape(-1)[source.to(positions_in.device)] + 1 + out["positions_out"] = expected + seq = None + if seq_lens is not None: + expected = seq_lens.clone() + if fill_seq_lens: + expected[:rows] = 1 + out["seq_lens"] = expected + seq = [int(v) for v in expected[:rows].tolist()] + kv = None + if kv_lens is not None: + expected = kv_lens.clone() + if kv_update == 1: + expected[:rows] += 1 + elif kv_update == 2: + expected[:rows] -= draft_len - rewind_accepted[:rows].to(expected.dtype) + out["kv_lens"] = expected + kv = [int(v) for v in expected[:rows].tolist()] + if req_idx_per_token is None: + return out + + num_tokens, lanes, tokens_per_block, data_bytes, scale_size, block_stride, scale_base = ( + dsa_layout + ) + req_idx = req_idx_per_token.clone() + fp8 = slot_mapping_fp8.clone() + scale = slot_mapping_scale.clone() + kv_indptr = gen_kv_indptr.clone() + cached_indptr = gen_cached_token_indptr.clone() + kv_indptr[0] = 0 + cached_indptr[0] = 0 + blocks = block_offsets.cpu().tolist() + max_blocks = block_offsets.shape[1] + flat_start = 0 + kv_sum = 0 + cached_sum = 0 + for r in range(rows): + kv_sum += kv[r] + cached_sum += kv[r] - seq[r] + kv_indptr[1 + r] = kv_sum + cached_indptr[1 + r] = cached_sum + start = kv[r] - seq[r] + for j in range(min(seq[r], lanes)): + g = flat_start + j + if g >= num_tokens: + break + position = start + j + # Truncating division, a non-negative remainder, then the clamp. + block = abs(position) // tokens_per_block * (1 if position >= 0 else -1) + offset = position - block * tokens_per_block + if offset < 0: + offset += tokens_per_block + block = min(max(block, 0), max_blocks - 1) + block_id = blocks[r][block] + req_idx[g] = r + fp8[g] = block_id * block_stride + offset * data_bytes + scale[g] = block_id * block_stride + scale_base + offset * scale_size + flat_start += seq[r] + out.update( + req_idx_per_token=req_idx, + slot_mapping_fp8=fp8, + slot_mapping_scale=scale, + gen_kv_indptr=kv_indptr, + gen_cached_token_indptr=cached_indptr, + ) + return out + + +def assert_state_args_valid( + rows: int, + draft_tokens: torch.Tensor, + positions_in: Optional[torch.Tensor], + positions_out: Optional[torch.Tensor], + gather_ids: Optional[torch.Tensor], + kv_update: int, + draft_len: int, + dsa_layout: List[int], + req_idx_per_token: Optional[torch.Tensor], +) -> None: + """The state arguments both draft-step ops accept and answer wrongly.""" + if positions_out is not None and positions_in is not None: + # Lanes read positions_in[gather_ids[r]] while others write positions_out[r]: overlapping buffers race. + lo_in, hi_in = positions_in.data_ptr(), positions_in.data_ptr() + positions_in.nbytes + lo_out, hi_out = positions_out.data_ptr(), positions_out.data_ptr() + positions_out.nbytes + assert hi_in <= lo_out or hi_out <= lo_in, "positions_out must not overlap positions_in" + if gather_ids is not None: + source = gather_ids[:rows] + assert bool(((source >= 0) & (source < positions_in.numel())).all()), ( + "gather_ids must index positions_in" + ) + # The rewind subtracts draft_len - accepted from each KV length: another draft length rewinds wrongly. + assert kv_update != 2 or draft_len == draft_tokens.shape[1], ( + "the KV rewind needs draft_len equal to the step's draft length" + ) + if req_idx_per_token is not None: + # No token lanes, or no tokens per block, writes no metadata or divides by zero. + assert len(dsa_layout) == 7 and dsa_layout[1] >= 1 and dsa_layout[2] >= 1, ( + "dsa_layout needs at least one token lane and one token per block" + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_token_table.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_token_table.py new file mode 100644 index 000000000000..d36ab203accd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_token_table.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The constant operands of ``glm_draft_ehproj_table``, built once from the checkpoint. + +eh_proj is row-parallel over ``cat([enorm(embed(t)), hnorm(h)], -1)``, so its output splits into +``W[:, :H] @ enorm(embed(t))``, which depends on the draft token only, and ``W[:, H:] @ hnorm(h)``. Each rank +tabulates the first half for the tokens of its LM-head vocabulary shard and keeps its output rows of the second. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul + +from ._draft_glue_rmsnorm import draft_glue_rmsnorm_cat + +#: Table rows normalized and multiplied per block, and output columns per block (bounds the FP32 copies). +_ROW_BLOCK = 2048 +_COL_BLOCK = 1536 + + +def build_draft_token_table( + eh_proj_weight: torch.Tensor, + embed_table: torch.Tensor, + enorm_weight: torch.Tensor, + eps: float, + tp_rank: int, + tp_size: int, + vocab_shard: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """``(token_table, hidden_rows)`` of rank ``tp_rank``. + + ``eh_proj_weight`` is the whole ``[H, 2H]`` BF16 eh_proj weight and ``embed_table`` the whole ``[V, H]`` BF16 + embedding. ``token_table[t] = bf16(W[:, :H] @ enorm(embed_table[tp_rank * vocab_shard + t]))``: enorm by the + draft-step glue kernel (the arithmetic of draft step 0's eh_proj input), the product in FP32 without TF32, + one rounding; a token past the table gives a zero row. ``hidden_rows`` is this rank's ``[H / tp_size, H]`` + slice ``W[tp_rank * H / tp : (tp_rank + 1) * H / tp, H:]``, the weight of the hidden-half GEMM. + """ + hidden = embed_table.shape[1] + share = hidden // tp_size + device = embed_table.device + hidden_rows = eh_proj_weight[tp_rank * share : (tp_rank + 1) * share, hidden:].contiguous() + token_table = torch.empty((vocab_shard, hidden), dtype=torch.bfloat16, device=device) + first = tp_rank * vocab_shard + ids = torch.arange(first, first + vocab_shard, dtype=torch.int64, device=device) + # The glue kernel's hidden-row operand is not read for a window inside the embedding half. + unused_rows = torch.empty((_ROW_BLOCK, hidden), dtype=torch.bfloat16, device=device) + with true_fp32_matmul(): + for row in range(0, vocab_shard, _ROW_BLOCK): + n = min(_ROW_BLOCK, vocab_shard - row) + normed = draft_glue_rmsnorm_cat( + embed_table, + ids[row : row + n], + unused_rows[:n], + enorm_weight, + enorm_weight, + eps, + 0, + hidden, + False, + ).float() + for lo in range(0, hidden, _COL_BLOCK): + weight_t = eh_proj_weight[lo : lo + _COL_BLOCK, :hidden].float().t() + token_table[row : row + n, lo : lo + _COL_BLOCK] = (normed @ weight_t).to( + torch.bfloat16 + ) + return token_table, hidden_rows diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_ehproj_table.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_ehproj_table.py new file mode 100644 index 000000000000..537bcf3460a2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_ehproj_table.py @@ -0,0 +1,284 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""An MTP draft step that also produces the next step's eh_proj output, via glm_draft_ehproj_table.""" + +from __future__ import annotations + +from typing import Dict, List, Optional, Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._draft_step_state import assert_state_args_valid, global_argmax, next_step_state + + +def _step(**overrides) -> dict: + spec = dict( + rows=1, + draft_len=7, + step=1, + vocab_shard=9680, + logits="random", + positions="plain", + fill=False, + kv_update=1, + dsa=True, + ) + spec.update(overrides) + return spec + + +class GlmDraftEhprojTable(OpWrapper): + """The greedy draft token of every row and the next step's eh_proj output with its all-reduce and norm. + + eh_proj is row-parallel over `cat([enorm(embed(t)), hnorm(h)], -1)`, so its output splits into + `W[:, :H] @ enorm(embed(t))`, which depends on the token only, and `W[:, H:] @ hnorm(h)`, which is known before + the draft LM head. `token_table` holds the first half for the tokens of this rank's LM-head shard + (`[vocab_shard, H]` BF16, one rounding of the FP32 product); `hidden_half` is this rank's FP32 share + `[rows, H / group_size]` of the second half (columns `[rank * H / tp, (rank + 1) * H / tp)`). + + Collective over the TP group, in one 8-CTA cluster per rank: each rank finds its logits shard's argmax and + publishes, through the MNNVL workspace `exchange_buffer` / `exchange_flags`, the (index, value) packet, the + token-table row of its candidate and its hidden-half share. Every rank then takes the global argmax in + torch.argmax's order into column `draft_step` of `draft_tokens`, writes + `residual_out = x = bf16(token_table_of_winner[token] + hidden_half)` (one FP32 add) and + `normed_out = bf16((x * rsqrt(mean(x^2) + eps)) * norm_weight)`: what eh_proj with its all-reduce, + input_layernorm and a zero residual returns. The next step's position ids, lengths and DSA decode metadata are + written as `glm_draft_sample_glue` writes them. + + Under PDL the kernel reads only the norm weight before its grid-dependency wait and triggers its dependents + right after it; the hidden-half GEMM on the other stream must be joined by an event before the launch. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # spec: rows (requests: 1 at c1, at most 2 where the target builds the table), draft_len (K), step, + # vocab_shard (154880 / 16 at tp16), logits ("random" or "tied"), positions ("gather" on step 0, else + # "plain"), fill, kv_update, dsa. At c1 the op serves draft steps 0-5: step 0 turns the 8 verified rows into + # one (gathered positions, seq_lens filled with 1, KV rewound by the rejected drafts), steps 1-5 advance by one. + CELLS: tuple[Cell, ...] = ( + Cell( + why="c1 draft step 0: gathered positions, seq_lens fill, KV rewind, DSA metadata", + spec=_step(step=0, positions="gather", fill=True, kv_update=2), + ), + Cell( + why="c1 draft steps 1-5: positions + 1, KV + 1, DSA metadata", + spec=_step(step=3), + ), + Cell( + why="c1 middle step, every rank's maximum equal: rank 0's row wins", + spec=_step(step=2, logits="tied"), + ), + Cell( + why="two requests, the most rows the target builds the table for", + spec=_step(rows=2, step=1), + ), + ) + + note = """ + Serves 2 to 16 ranks whose shares of the 6144 hidden columns are whole 8-column vectors; at tp16 a CTA's 768 + columns span two ranks' 384-column shares, a mapping a 4-rank run does not reach. The workspace comes from + `DraftExchangeWorkspace` (one per op, built outside capture). + + `normed_out` differs from eh_proj + all-reduce + RMSNorm in the order of the norm's sum of squares (CTA + slices in order); `residual_out` differs from eh_proj's tp partial sums rounded to BF16 before their reduction: + here only the token half is rounded before x. Both feed the draft only. + """ + + def __call__( + self, + logits: torch.Tensor, + exchange_buffer: torch.Tensor, + exchange_flags: torch.Tensor, + group_size: int, + draft_tokens: torch.Tensor, + draft_step: int, + hidden_half: torch.Tensor, + token_table: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + normed_out: torch.Tensor, + residual_out: torch.Tensor, + positions_in: Optional[torch.Tensor], + positions_out: Optional[torch.Tensor], + gather_ids: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + fill_seq_lens: bool, + kv_lens: Optional[torch.Tensor], + kv_update: int, + rewind_accepted: Optional[torch.Tensor], + draft_len: int, + block_offsets: Optional[torch.Tensor], + dsa_layout: List[int], + req_idx_per_token: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + gen_kv_indptr: Optional[torch.Tensor], + gen_cached_token_indptr: Optional[torch.Tensor], + ) -> None: + torch.ops.trtllm.glm_draft_ehproj_table( + logits, + exchange_buffer, + exchange_flags, + group_size, + draft_tokens, + draft_step, + hidden_half, + token_table, + norm_weight, + eps, + normed_out, + residual_out, + positions_in, + positions_out, + gather_ids, + seq_lens, + fill_seq_lens, + kv_lens, + kv_update, + rewind_accepted, + draft_len, + block_offsets, + dsa_layout, + req_idx_per_token, + slot_mapping_fp8, + slot_mapping_scale, + gen_kv_indptr, + gen_cached_token_indptr, + ) + + def reference( + self, + logits: Sequence[torch.Tensor], + exchange_buffer: torch.Tensor, + exchange_flags: torch.Tensor, + group_size: int, + draft_tokens: torch.Tensor, + draft_step: int, + hidden_half: Sequence[torch.Tensor], + token_table: Sequence[torch.Tensor], + norm_weight: torch.Tensor, + eps: float, + normed_out: torch.Tensor, + residual_out: torch.Tensor, + positions_in: Optional[torch.Tensor], + positions_out: Optional[torch.Tensor], + gather_ids: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + fill_seq_lens: bool, + kv_lens: Optional[torch.Tensor], + kv_update: int, + rewind_accepted: Optional[torch.Tensor], + draft_len: int, + block_offsets: Optional[torch.Tensor], + dsa_layout: List[int], + req_idx_per_token: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + gen_kv_indptr: Optional[torch.Tensor], + gen_cached_token_indptr: Optional[torch.Tensor], + ) -> Dict[str, torch.Tensor]: + """What the op leaves in each buffer it writes, by argument name. + + The documented exception to mirroring `__call__`: `logits`, `hidden_half` and `token_table` are the + group's, in rank order, because every rank's output depends on all of them. Every other argument is the + calling rank's, read as it was before the call. The exchange workspace is not consulted. + """ + rows = logits[0].shape[0] + vocab_shard = logits[0].shape[1] + tokens = global_argmax(logits) + expected_draft = draft_tokens.clone() + expected_draft[:rows, draft_step] = tokens + token_half = torch.stack( + [ + token_table[int(t) // vocab_shard][int(t) % vocab_shard].float() + for t in tokens.tolist() + ] + ) + x = (token_half + torch.cat([h.float() for h in hidden_half], dim=-1)).to(torch.bfloat16) + rstd = torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + eps) + normed = ((x.float() * rstd) * norm_weight.float()).to(torch.bfloat16) + out = {"draft_tokens": expected_draft, "residual_out": x, "normed_out": normed} + out.update( + next_step_state( + rows, + positions_in, + positions_out, + gather_ids, + seq_lens, + fill_seq_lens, + kv_lens, + kv_update, + rewind_accepted, + draft_len, + block_offsets, + dsa_layout, + req_idx_per_token, + slot_mapping_fp8, + slot_mapping_scale, + gen_kv_indptr, + gen_cached_token_indptr, + ) + ) + return out + + def is_valid( + self, + logits: torch.Tensor, + exchange_buffer: torch.Tensor, + exchange_flags: torch.Tensor, + group_size: int, + draft_tokens: torch.Tensor, + draft_step: int, + hidden_half: torch.Tensor, + token_table: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + normed_out: torch.Tensor, + residual_out: torch.Tensor, + positions_in: Optional[torch.Tensor], + positions_out: Optional[torch.Tensor], + gather_ids: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + fill_seq_lens: bool, + kv_lens: Optional[torch.Tensor], + kv_update: int, + rewind_accepted: Optional[torch.Tensor], + draft_len: int, + block_offsets: Optional[torch.Tensor], + dsa_layout: List[int], + req_idx_per_token: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + gen_kv_indptr: Optional[torch.Tensor], + gen_cached_token_indptr: Optional[torch.Tensor], + ) -> None: + assert_state_args_valid( + logits.shape[0], + draft_tokens, + positions_in, + positions_out, + gather_ids, + kv_update, + draft_len, + dsa_layout, + req_idx_per_token, + ) + # x is stored row by row while the cluster still reads the token table for the later rows, and the norm + # after the last one: an output overlapping an input, or the other output, replaces values the kernel has + # yet to read or to return. + def overlaps(a: torch.Tensor, b: torch.Tensor) -> bool: + return a.data_ptr() < b.data_ptr() + b.nbytes and b.data_ptr() < a.data_ptr() + a.nbytes + + for name, t in (("normed_out", normed_out), ("residual_out", residual_out)): + assert not any(overlaps(t, other) for other in (logits, hidden_half, token_table)), ( + f"{name} must not overlap the op's inputs" + ) + assert not overlaps(normed_out, residual_out), ( + "normed_out and residual_out must be separate buffers" + ) + + +glm_draft_ehproj_table = GlmDraftEhprojTable() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_sample_glue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_sample_glue.py new file mode 100644 index 000000000000..741c387dedcf --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_sample_glue.py @@ -0,0 +1,322 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One MTP draft step after the vocabulary-sharded draft LM head, in one kernel, via glm_draft_sample_glue.""" + +from __future__ import annotations + +from typing import Dict, List, Optional, Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._draft_step_state import ( + assert_state_args_valid, + eh_proj_input_window, + global_argmax, + next_step_state, +) + + +def _step(**overrides) -> dict: + spec = dict( + rows=1, + draft_len=7, + step=6, + vocab_shard=9680, + logits="random", + glue=False, + positions=None, + fill=False, + kv_update=1, + dsa=True, + next_new=False, + ) + spec.update(overrides) + return spec + + +class GlmDraftSampleGlue(OpWrapper): + """The greedy draft token of every row, and optionally the next step's inputs and state, in one CTA. + + Collective over the TP group (`group_size` ranks, the workspace's group): each rank takes the argmax of its + logits shard `[rows, vocab_shard]` (rank r holds tokens `[r * vocab_shard, (r + 1) * vocab_shard)`), the ranks + exchange the (index, value) pairs through the MNNVL workspace `exchange_buffer` / `exchange_flags`, and every + rank writes the global argmax, in torch.argmax's order, into column `draft_step` of `draft_tokens`. Then, each + part skipped when its output is None: + - `glue_out`: this rank's window `[col_lo, col_hi)` of `cat([enorm(embed_table[token]), hnorm(hidden)], -1)`, + the next step's eh_proj input (FlashInfer's CuTe RMSNorm arithmetic), contiguous; + - `positions_out[r] = positions_in[gather_ids[r] if gather_ids is not None else r] + 1`; + - `seq_lens` filled with 1 when `fill_seq_lens`; `kv_lens` advanced as `kv_update` says (0 none, 1 `+ 1`, + 2 `- (draft_len - rewind_accepted)`); + - the DSA decode metadata of the next step from those lengths (`req_idx_per_token`, the two slot mappings, the + two generation indptrs; `dsa_layout` = [tokens, token lanes per request, tokens per block, data bytes per + token, scale bytes per token, bytes per block, scale offset in a block]); + - `next_new_tokens[r] = [accepted_tokens[batch_indices[r], num_accepted[r] - 1], draft tokens...]`, the + output rows after the last draft step. + + Under PDL the kernel reads only the norm weights before its grid-dependency wait and triggers its dependents + right after it. The workspace serves this op only (each call empties the buffer of the previous one). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # spec: rows (requests), draft_len (K), step, vocab_shard (the per-rank LM-head shard: 154880 / 16 at tp16), + # logits ("random", or "tied": every rank's maximum equal), glue (the eh_proj input window, this rank's + # 2 * 6144 / tp columns), positions (None, "gather" or "plain"), fill (seq_lens), kv_update, dsa (metadata), + # next_new (output rows). At c1 the token table serves draft steps 0-5 and this op the last one (step 6): + # token, KV + 1, DSA metadata and the next-new-tokens rows. Above two requests per rank (no token table) it + # serves every step, with the eh_proj input window. + CELLS: tuple[Cell, ...] = ( + Cell( + why="c1 last draft step: token, KV + 1, DSA metadata, next-new-tokens rows", + spec=_step(next_new=True), + ), + Cell( + why="c1 last draft step, every rank's maximum equal: rank 0's index wins", + spec=_step(logits="tied", next_new=True), + ), + Cell( + why="step 0 without the token table: window, gathered positions, fill, rewind", + spec=_step(rows=4, step=0, glue=True, positions="gather", fill=True, kv_update=2), + ), + Cell( + why="a middle step without the token table: window, positions + 1, KV + 1", + spec=_step(rows=4, step=3, glue=True, positions="plain"), + ), + Cell( + why="the most rows one launch serves (8 requests), last step", + spec=_step(rows=8, next_new=True), + ), + ) + + note = """ + The workspace comes from `DraftExchangeWorkspace` (one per op, built outside capture). Ranks must call in + lockstep with the same `rows`; a missing rank leaves the others polling. + + The vocabulary shards must be equal and in rank order, which an unpadded vocabulary-parallel LM head gives + (154880 = 16 x 9680). + """ + + def __call__( + self, + logits: torch.Tensor, + exchange_buffer: torch.Tensor, + exchange_flags: torch.Tensor, + group_size: int, + draft_tokens: torch.Tensor, + draft_step: int, + embed_table: Optional[torch.Tensor], + hidden: Optional[torch.Tensor], + enorm_weight: Optional[torch.Tensor], + hnorm_weight: Optional[torch.Tensor], + eps: float, + col_lo: int, + col_hi: int, + glue_out: Optional[torch.Tensor], + positions_in: Optional[torch.Tensor], + positions_out: Optional[torch.Tensor], + gather_ids: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + fill_seq_lens: bool, + kv_lens: Optional[torch.Tensor], + kv_update: int, + rewind_accepted: Optional[torch.Tensor], + draft_len: int, + block_offsets: Optional[torch.Tensor], + dsa_layout: List[int], + req_idx_per_token: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + gen_kv_indptr: Optional[torch.Tensor], + gen_cached_token_indptr: Optional[torch.Tensor], + accepted_tokens: Optional[torch.Tensor], + batch_indices: Optional[torch.Tensor], + num_accepted: Optional[torch.Tensor], + next_new_tokens: Optional[torch.Tensor], + ) -> None: + torch.ops.trtllm.glm_draft_sample_glue( + logits, + exchange_buffer, + exchange_flags, + group_size, + draft_tokens, + draft_step, + embed_table, + hidden, + enorm_weight, + hnorm_weight, + eps, + col_lo, + col_hi, + glue_out, + positions_in, + positions_out, + gather_ids, + seq_lens, + fill_seq_lens, + kv_lens, + kv_update, + rewind_accepted, + draft_len, + block_offsets, + dsa_layout, + req_idx_per_token, + slot_mapping_fp8, + slot_mapping_scale, + gen_kv_indptr, + gen_cached_token_indptr, + accepted_tokens, + batch_indices, + num_accepted, + next_new_tokens, + ) + + def reference( + self, + logits: Sequence[torch.Tensor], + exchange_buffer: torch.Tensor, + exchange_flags: torch.Tensor, + group_size: int, + draft_tokens: torch.Tensor, + draft_step: int, + embed_table: Optional[torch.Tensor], + hidden: Optional[torch.Tensor], + enorm_weight: Optional[torch.Tensor], + hnorm_weight: Optional[torch.Tensor], + eps: float, + col_lo: int, + col_hi: int, + glue_out: Optional[torch.Tensor], + positions_in: Optional[torch.Tensor], + positions_out: Optional[torch.Tensor], + gather_ids: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + fill_seq_lens: bool, + kv_lens: Optional[torch.Tensor], + kv_update: int, + rewind_accepted: Optional[torch.Tensor], + draft_len: int, + block_offsets: Optional[torch.Tensor], + dsa_layout: List[int], + req_idx_per_token: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + gen_kv_indptr: Optional[torch.Tensor], + gen_cached_token_indptr: Optional[torch.Tensor], + accepted_tokens: Optional[torch.Tensor], + batch_indices: Optional[torch.Tensor], + num_accepted: Optional[torch.Tensor], + next_new_tokens: Optional[torch.Tensor], + ) -> Dict[str, torch.Tensor]: + """What the op leaves in each buffer it writes, by argument name. + + The documented exception to mirroring `__call__`: the token is a function of every rank's logits, so + `logits` is the group's shards in rank order. Every other argument is the calling rank's, and the buffers + the op updates in place are read as they were before the call. The exchange workspace is not consulted. + """ + rows = logits[0].shape[0] + tokens = global_argmax(logits) + expected_draft = draft_tokens.clone() + expected_draft[:rows, draft_step] = tokens + out = {"draft_tokens": expected_draft} + if glue_out is not None: + out["glue_out"] = eh_proj_input_window( + tokens, embed_table, hidden, enorm_weight, hnorm_weight, eps, col_lo, col_hi + ) + out.update( + next_step_state( + rows, + positions_in, + positions_out, + gather_ids, + seq_lens, + fill_seq_lens, + kv_lens, + kv_update, + rewind_accepted, + draft_len, + block_offsets, + dsa_layout, + req_idx_per_token, + slot_mapping_fp8, + slot_mapping_scale, + gen_kv_indptr, + gen_cached_token_indptr, + ) + ) + if next_new_tokens is not None: + expected = next_new_tokens.clone().view(rows, -1) + last = accepted_tokens[batch_indices[:rows].long(), num_accepted[:rows].long() - 1] + expected[:, 0] = last + expected[:, 1:] = expected_draft[:rows] + out["next_new_tokens"] = expected.view_as(next_new_tokens) + return out + + def is_valid( + self, + logits: torch.Tensor, + exchange_buffer: torch.Tensor, + exchange_flags: torch.Tensor, + group_size: int, + draft_tokens: torch.Tensor, + draft_step: int, + embed_table: Optional[torch.Tensor], + hidden: Optional[torch.Tensor], + enorm_weight: Optional[torch.Tensor], + hnorm_weight: Optional[torch.Tensor], + eps: float, + col_lo: int, + col_hi: int, + glue_out: Optional[torch.Tensor], + positions_in: Optional[torch.Tensor], + positions_out: Optional[torch.Tensor], + gather_ids: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + fill_seq_lens: bool, + kv_lens: Optional[torch.Tensor], + kv_update: int, + rewind_accepted: Optional[torch.Tensor], + draft_len: int, + block_offsets: Optional[torch.Tensor], + dsa_layout: List[int], + req_idx_per_token: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + gen_kv_indptr: Optional[torch.Tensor], + gen_cached_token_indptr: Optional[torch.Tensor], + accepted_tokens: Optional[torch.Tensor], + batch_indices: Optional[torch.Tensor], + num_accepted: Optional[torch.Tensor], + next_new_tokens: Optional[torch.Tensor], + ) -> None: + rows = logits.shape[0] + assert_state_args_valid( + rows, + draft_tokens, + positions_in, + positions_out, + gather_ids, + kv_update, + draft_len, + dsa_layout, + req_idx_per_token, + ) + if next_new_tokens is not None: + # The kernel writes row r at offset r * (draft_len + 1) and reads + # accepted_tokens[batch_indices[r], num_accepted[r] - 1] unbounded. + assert tuple(next_new_tokens.shape) == (rows, draft_tokens.shape[1] + 1), ( + "next_new_tokens must be exactly [rows, draft_len + 1]" + ) + index = batch_indices[:rows].long() + count = num_accepted[:rows].long() + assert bool(((index >= 0) & (index < accepted_tokens.shape[0])).all()), ( + "batch_indices must index accepted_tokens" + ) + assert bool(((count >= 1) & (count <= accepted_tokens.shape[1])).all()), ( + "num_accepted must lie in [1, accepted_tokens.shape[1]]" + ) + + +glm_draft_sample_glue = GlmDraftSampleGlue() diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index bf24256133dd..563268303c40 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -949,6 +949,26 @@ def _(target_tokens: torch.Tensor, draft_tokens: torch.Tensor, return draft_tokens.new_empty( (num_contexts + draft_tokens.shape[0], )) + @torch.library.register_fake("trtllm::glm_draft_sample_glue") + def _(logits, exchange_buffer, exchange_flags, group_size, draft_tokens, + draft_step, embed_table, hidden, enorm_weight, hnorm_weight, eps, + col_lo, col_hi, glue_out, positions_in, positions_out, gather_ids, + seq_lens, fill_seq_lens, kv_lens, kv_update, rewind_accepted, + draft_len, block_offsets, dsa_layout, req_idx_per_token, + slot_mapping_fp8, slot_mapping_scale, gen_kv_indptr, + gen_cached_token_indptr, accepted_tokens, batch_indices, + num_accepted, next_new_tokens): + return None + + @torch.library.register_fake("trtllm::glm_draft_ehproj_table") + def _(logits, exchange_buffer, exchange_flags, group_size, draft_tokens, + draft_step, hidden_half, token_table, norm_weight, eps, normed_out, + residual_out, positions_in, positions_out, gather_ids, seq_lens, + fill_seq_lens, kv_lens, kv_update, rewind_accepted, draft_len, + block_offsets, dsa_layout, req_idx_per_token, slot_mapping_fp8, + slot_mapping_scale, gen_kv_indptr, gen_cached_token_indptr): + return None + @torch.library.register_fake("trtllm::fused_compute_probs_from_logits") def _( logits: torch.Tensor, diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_op_matrix.py new file mode 100644 index 000000000000..2dcdc782c683 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_op_matrix.py @@ -0,0 +1,627 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for the two collective MTP draft-step entries of the GLM-5.2 tp16 target. + +`glm_draft_ehproj_table` and `glm_draft_sample_glue` exchange through an MNNVL workspace across the TP group, +so neither can run in one process. Run plainly, this script is its own launcher: it spawns one rank per device +named in CUDA_VISIBLE_DEVICES, under a deadline it enforces by killing every rank's session, because a broken +exchange hangs (a rank polls for a packet that never lands) rather than raising. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _draft_step_op_matrix.py draft_ehproj_table + +Each rank drives every cell of the entry, a draft loop of the target's shape (six table steps and the last glue +step, repeated) eagerly and as CUDA-graph replays, and the entry's guards. The 16-rank geometry of the table +kernel (a CTA's columns spanning two ranks' hidden-half shares) needs 16 ranks over 4 nodes: start one process +per GPU with an external launcher and `--rank-worker`, with RANK, WORLD_SIZE, LOCAL_RANK, MASTER_ADDR and +MASTER_PORT in the environment (SLURM_PROCID, SLURM_NTASKS and SLURM_LOCALID stand in for the first three): + + srun -N4 --ntasks-per-node=4 python _draft_step_op_matrix.py draft_ehproj_table --rank-worker + +Not a pytest module: the collected entry points are `test_modeling_v2_draft_ehproj_table.py` and +`test_modeling_v2_draft_sample_glue.py` beside this file. The launcher half imports nothing from +tensorrt_llm, and the ranks reach the target's ops by absolute import. +""" + +import os +import signal +import socket +import subprocess +import sys +import time +import zlib +from pathlib import Path +from typing import Any, Dict, List + +import torch +import torch.distributed as dist + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 1500 + +HIDDEN = 6144 +EPS = 1e-5 +MAX_BLOCKS = 64 +#: The DSA indexer-K cache layout of GLM-5.2: 64 tokens per block, 128 FP8 bytes and one FP32 scale per token. +TOKENS_PER_BLOCK = 64 +DATA_BYTES = 128 +SCALE_BYTES = 4 +TOKEN_LANES = 8 + +# Bound in the rank body: importing the entries pulls in tensorrt_llm. +RANK = 0 +WORLD = 1 +HOST_GROUP: Any = None +table_op: Any = None +glue_op: Any = None +validating: Any = None +assert_within_ulp: Any = None +WORKSPACES: Dict[str, Any] = {} + + +def _generator(*key: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(zlib.crc32(repr(key).encode()) % 2**31) + + +def _logits(spec: dict, seed: int, rank: int) -> torch.Tensor: + """Rank `rank`'s BF16 logits shard; every rank can rebuild every other rank's.""" + rows, shard = spec["rows"], spec["vocab_shard"] + g = _generator(seed, 1, rank) + logits = torch.randn(rows, shard, generator=g, device="cuda") + if spec["logits"] == "tied": + # Every rank's maximum is 8 at two local positions: the lower position of rank 0 wins. + first = torch.randint(0, shard // 2, (rows,), generator=g, device="cuda") + second = first + shard // 2 + rows_index = torch.arange(rows, device="cuda") + logits[rows_index, first] = 8.0 + logits[rows_index, second] = 8.0 + return logits.to(torch.bfloat16) + + +class _LazyTables: + """Every rank's token table, built on first use: the reference reads only the winners' rows.""" + + def __init__(self, seed: int, shard: int): + self._seed, self._shard, self._tables = seed, shard, {} + + def __getitem__(self, rank: int) -> torch.Tensor: + if rank not in self._tables: + g = _generator(self._seed, 2, rank) + self._tables[rank] = torch.randn( + self._shard, HIDDEN, generator=g, device="cuda", dtype=torch.bfloat16 + ) + return self._tables[rank] + + def __len__(self) -> int: + return WORLD + + +def _hidden_half(seed: int, rank: int, rows: int) -> torch.Tensor: + g = _generator(seed, 3, rank) + return torch.randn(rows, HIDDEN // WORLD, generator=g, device="cuda") + + +def _state(spec: dict, seed: int) -> Dict[str, Any]: + """The step's state arguments, identical on every rank.""" + rows, draft_len, step = spec["rows"], spec["draft_len"], spec["step"] + g = _generator(seed, 4) + vocab = spec["vocab_shard"] * WORLD + num_accepted = torch.randint(1, draft_len + 2, (rows,), generator=g, device="cuda").int() + batch_indices = torch.randperm(rows, generator=g, device="cuda").int() + state: Dict[str, Any] = dict( + draft_tokens=torch.randint(0, vocab, (rows, draft_len), generator=g, device="cuda").int(), + draft_step=step, + positions_in=None, + positions_out=None, + gather_ids=None, + seq_lens=None, + fill_seq_lens=spec["fill"], + kv_lens=None, + kv_update=spec["kv_update"], + rewind_accepted=None, + draft_len=draft_len, + block_offsets=None, + dsa_layout=[], + req_idx_per_token=None, + slot_mapping_fp8=None, + slot_mapping_scale=None, + gen_kv_indptr=None, + gen_cached_token_indptr=None, + ) + verify_rows = rows * (draft_len + 1) + if spec["positions"] == "gather": + state["positions_in"] = torch.randint( + 100, 100000, (verify_rows,), generator=g, device="cuda" + ).int() + state["gather_ids"] = batch_indices.long() * (draft_len + 1) + num_accepted.long() - 1 + state["positions_out"] = torch.full((rows,), -9, dtype=torch.int32, device="cuda") + elif spec["positions"] == "plain": + state["positions_in"] = torch.randint( + 100, 100000, (rows,), generator=g, device="cuda" + ).int() + state["positions_out"] = torch.full((rows,), -9, dtype=torch.int32, device="cuda") + seq_before = draft_len + 1 if spec["fill"] else 1 + state["seq_lens"] = torch.full((rows,), seq_before, dtype=torch.int32, device="cuda") + state["kv_lens"] = ( + 1000 + 37 * torch.arange(rows, device="cuda", dtype=torch.int32) + 3 * step + ).int() + if spec["kv_update"] == 2: + state["rewind_accepted"] = num_accepted + if spec["dsa"]: + num_tokens = rows + state["block_offsets"] = torch.randint( + 0, 4096, (rows, MAX_BLOCKS), generator=g, device="cuda" + ).int() + state["dsa_layout"] = [ + num_tokens, + TOKEN_LANES, + TOKENS_PER_BLOCK, + DATA_BYTES, + SCALE_BYTES, + TOKENS_PER_BLOCK * (DATA_BYTES + SCALE_BYTES), + TOKENS_PER_BLOCK * DATA_BYTES, + ] + state["req_idx_per_token"] = torch.full((num_tokens + 3,), -5, dtype=torch.int32).cuda() + state["slot_mapping_fp8"] = torch.full((num_tokens + 3,), -5, dtype=torch.int64).cuda() + state["slot_mapping_scale"] = torch.full((num_tokens + 3,), -5, dtype=torch.int64).cuda() + state["gen_kv_indptr"] = torch.full((rows + 4,), -5, dtype=torch.int64).cuda() + state["gen_cached_token_indptr"] = torch.full((rows + 4,), -5, dtype=torch.int64).cuda() + state["_num_accepted"] = num_accepted + state["_batch_indices"] = batch_indices + return state + + +def _snapshot(state: Dict[str, Any]) -> Dict[str, Any]: + return { + k: (v.clone() if isinstance(v, torch.Tensor) else v) + for k, v in state.items() + if not k.startswith("_") + } + + +def _assert_exact(name: str, out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + try: + torch.testing.assert_close(out, ref, rtol=0, atol=0) + except AssertionError as exc: + raise AssertionError(f"{where}: {name}: {exc}") from exc + + +def _compare( + expected: Dict[str, torch.Tensor], actual: Dict[str, torch.Tensor], where: str +) -> None: + """Integers and x bit for bit; the normalized rows within one RMSNorm's rounding. + + The RMSNorm outputs (`normed_out`, `glue_out`) differ from the FP32 reference in the order of the 6144 + squares and in the kernel's approximate reciprocal square root: a few 1e-7 relative on the scale, one BF16 + ulp on the rare element that close to a rounding midpoint. So 1 ulp against the row scale and 0.25 ulp of + relative RMS, as for the step-0 glue kernel. + """ + for name, ref in expected.items(): + out = actual[name] + if name in ("normed_out", "glue_out"): + try: + assert_within_ulp(out, ref, element_ulp=1.0, rms_ulp=0.25) + except AssertionError as exc: + raise AssertionError(f"{where}: {name}: {exc}") from exc + else: + _assert_exact(name, out, ref, where) + + +def _table_call(spec: dict, seed: int, logits: List[torch.Tensor], tables: _LazyTables): + """One table-op call on this rank: (call kwargs, reference kwargs).""" + rows = spec["rows"] + state = _state(spec, seed) + halves = [_hidden_half(seed, r, rows) for r in range(WORLD)] + norm_g = _generator(seed, 5) + norm_weight = (1.0 + 0.1 * torch.randn(HIDDEN, generator=norm_g, device="cuda")).to( + torch.bfloat16 + ) + workspace = WORKSPACES["table"] + kwargs = dict( + logits=logits[RANK], + exchange_buffer=workspace.buffer, + exchange_flags=workspace.flags, + group_size=WORLD, + hidden_half=halves[RANK], + token_table=tables[RANK], + norm_weight=norm_weight, + eps=EPS, + normed_out=torch.full((rows, HIDDEN), 7.0, dtype=torch.bfloat16, device="cuda"), + residual_out=torch.full((rows, HIDDEN), 7.0, dtype=torch.bfloat16, device="cuda"), + **{k: v for k, v in state.items() if not k.startswith("_")}, + ) + ref_kwargs = dict( + kwargs, **_snapshot(state), logits=logits, hidden_half=halves, token_table=tables + ) + return kwargs, ref_kwargs + + +def _glue_call(spec: dict, seed: int, logits: List[torch.Tensor], embed_table: torch.Tensor): + """One glue-op call on this rank: (call kwargs, reference kwargs).""" + rows, draft_len = spec["rows"], spec["draft_len"] + state = _state(spec, seed) + g = _generator(seed, 6) + window = 2 * HIDDEN // WORLD + kwargs = dict( + logits=logits[RANK], + exchange_buffer=WORKSPACES["glue"].buffer, + exchange_flags=WORKSPACES["glue"].flags, + group_size=WORLD, + embed_table=None, + hidden=None, + enorm_weight=None, + hnorm_weight=None, + eps=EPS, + col_lo=0, + col_hi=0, + glue_out=None, + accepted_tokens=None, + batch_indices=None, + num_accepted=None, + next_new_tokens=None, + **{k: v for k, v in state.items() if not k.startswith("_")}, + ) + if spec["glue"]: + kwargs.update( + embed_table=embed_table, + hidden=(3.0 * torch.randn(rows, HIDDEN, generator=g, device="cuda")).to(torch.bfloat16), + enorm_weight=(1.0 + 0.1 * torch.randn(HIDDEN, generator=g, device="cuda")).to( + torch.bfloat16 + ), + hnorm_weight=(1.0 + 0.1 * torch.randn(HIDDEN, generator=g, device="cuda")).to( + torch.bfloat16 + ), + col_lo=RANK * window, + col_hi=(RANK + 1) * window, + glue_out=torch.full((rows, window), 7.0, dtype=torch.bfloat16, device="cuda"), + ) + if spec["next_new"]: + kwargs.update( + accepted_tokens=torch.randint( + 0, 1000, (rows, draft_len + 1), generator=g, device="cuda" + ).int(), + batch_indices=state["_batch_indices"], + num_accepted=state["_num_accepted"], + next_new_tokens=torch.full((rows, draft_len + 1), -3, dtype=torch.int32, device="cuda"), + ) + ref_kwargs = dict(kwargs, **_snapshot(state), logits=logits) + return kwargs, ref_kwargs + + +_TABLE_OUTPUTS = ("draft_tokens", "normed_out", "residual_out") +_GLUE_OUTPUTS = ("draft_tokens", "glue_out", "next_new_tokens") +_STATE_OUTPUTS = ( + "positions_out", + "seq_lens", + "kv_lens", + "req_idx_per_token", + "slot_mapping_fp8", + "slot_mapping_scale", + "gen_kv_indptr", + "gen_cached_token_indptr", +) + + +def _written(kwargs: Dict[str, Any], names) -> Dict[str, torch.Tensor]: + return {n: kwargs[n] for n in names if kwargs.get(n) is not None} + + +def check_table_cells() -> int: + for index, cell in enumerate(table_op.CELLS): + spec = cell.spec + seed = zlib.crc32(cell.why.encode()) % 2**31 + logits = [_logits(spec, seed, r) for r in range(WORLD)] + tables = _LazyTables(seed, spec["vocab_shard"]) + kwargs, ref_kwargs = _table_call(spec, seed, logits, tables) + with validating(table_op): + table_op(**kwargs) + torch.cuda.synchronize() + expected = table_op.reference(**ref_kwargs) + _compare(expected, _written(kwargs, _TABLE_OUTPUTS + _STATE_OUTPUTS), f"cell {index}") + return len(table_op.CELLS) + + +def check_glue_cells() -> int: + embed_table = None + for index, cell in enumerate(glue_op.CELLS): + spec = cell.spec + seed = zlib.crc32(cell.why.encode()) % 2**31 + if spec["glue"] and embed_table is None: + embed_table = torch.randn( + spec["vocab_shard"] * WORLD, + HIDDEN, + generator=_generator(9), + device="cuda", + dtype=torch.bfloat16, + ) + logits = [_logits(spec, seed, r) for r in range(WORLD)] + kwargs, ref_kwargs = _glue_call(spec, seed, logits, embed_table) + with validating(glue_op): + glue_op(**kwargs) + torch.cuda.synchronize() + expected = glue_op.reference(**ref_kwargs) + _compare(expected, _written(kwargs, _GLUE_OUTPUTS + _STATE_OUTPUTS), f"cell {index}") + return len(glue_op.CELLS) + + +def _draft_loop_specs(draft_len: int = 7): + """The target's c1 draft loop: table steps 0..K-2, then the glue for step K-1.""" + table_steps = [] + for step in range(draft_len - 1): + table_steps.append( + table_op.CELLS[0].spec + if step == 0 + else dict(table_op.CELLS[1].spec, step=step) + ) + glue_step = dict(glue_op.CELLS[0].spec, step=draft_len - 1) + return table_steps, glue_step + + +def check_draft_loop_eager_and_graph() -> None: + """Seven consecutive calls per decode step over two workspaces, eagerly and as graph replays. + + The Lamport flags rotate on every call and each call empties the buffer of the previous one, so the + exchange must stay correct across steps, decode steps and replays of a captured loop. + """ + table_steps, glue_step = _draft_loop_specs() + calls = [] + for step, spec in enumerate(table_steps): + seed = 1000 + step + logits = [_logits(spec, seed, r) for r in range(WORLD)] + tables = _LazyTables(seed, spec["vocab_shard"]) + calls.append(("table", _table_call(spec, seed, logits, tables))) + seed = 1000 + len(table_steps) + logits = [_logits(glue_step, seed, r) for r in range(WORLD)] + calls.append(("glue", _glue_call(glue_step, seed, logits, None))) + + def run_loop() -> None: + for kind, (kwargs, _) in calls: + (table_op if kind == "table" else glue_op)(**kwargs) + + def reset_and_expect(): + """Restore every buffer the loop updates in place; return the expected outputs per call.""" + expected = [] + for kind, (kwargs, ref_kwargs) in calls: + for name, value in ref_kwargs.items(): + # The snapshots of the buffers the op updates in place are separate tensors. + live = kwargs.get(name) + if isinstance(value, torch.Tensor) and isinstance(live, torch.Tensor): + if live is not value: + live.copy_(value) + op = table_op if kind == "table" else glue_op + expected.append(op.reference(**ref_kwargs)) + return expected + + for repeat in range(3): + expected = reset_and_expect() + run_loop() + torch.cuda.synchronize() + for (kind, (kwargs, _)), exp in zip(calls, expected): + names = (_TABLE_OUTPUTS if kind == "table" else _GLUE_OUTPUTS) + _STATE_OUTPUTS + _compare(exp, _written(kwargs, names), f"eager loop {repeat} {kind}") + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(side): + with torch.cuda.graph(graph, stream=side): + run_loop() + torch.cuda.current_stream().wait_stream(side) + for repeat in range(3): + expected = reset_and_expect() + graph.replay() + torch.cuda.synchronize() + for (kind, (kwargs, _)), exp in zip(calls, expected): + names = (_TABLE_OUTPUTS if kind == "table" else _GLUE_OUTPUTS) + _STATE_OUTPUTS + _compare(exp, _written(kwargs, names), f"graph replay {repeat} {kind}") + + +def check_guards_refuse() -> None: + """The guards refuse their silent-wrong inputs on every rank before anything is launched.""" + spec = dict(table_op.CELLS[0].spec) + seed = 77 + logits = [_logits(spec, seed, r) for r in range(WORLD)] + kwargs, _ = _table_call(spec, seed, logits, _LazyTables(seed, spec["vocab_shard"])) + bad = dict(kwargs, positions_out=kwargs["positions_in"]) + try: + with validating(table_op): + table_op(**bad) + except AssertionError as exc: + assert "must not overlap" in str(exc), exc + else: + raise AssertionError("the table guard accepted positions_out aliasing positions_in") + bad = dict(kwargs, draft_len=kwargs["draft_len"] + 1) + try: + with validating(table_op): + table_op(**bad) + except AssertionError as exc: + assert "KV rewind" in str(exc), exc + else: + raise AssertionError("the table guard accepted a rewind with another draft length") + spec = dict(glue_op.CELLS[0].spec) + logits = [_logits(spec, seed, r) for r in range(WORLD)] + kwargs, _ = _glue_call(spec, seed, logits, None) + wide = torch.zeros( + spec["rows"], spec["draft_len"] + 2, dtype=torch.int32, device="cuda" + ) + try: + with validating(glue_op): + glue_op(**dict(kwargs, next_new_tokens=wide)) + except AssertionError as exc: + assert "next_new_tokens must be exactly" in str(exc), exc + else: + raise AssertionError("the glue guard accepted a wider next_new_tokens") + + +def _rank_env(): + env = os.environ + rank = int(env.get("RANK", env.get("SLURM_PROCID", "0"))) + world = int(env.get("WORLD_SIZE", env.get("SLURM_NTASKS", "1"))) + local_rank = int(env.get("LOCAL_RANK", env.get("SLURM_LOCALID", str(rank)))) + return rank, world, local_rank + + +def _run_one_rank(entry: str, expected_cells: int) -> int: + """Body of one rank: build the workspaces, run the entry's checks, exit hard on a failure.""" + global RANK, WORLD, HOST_GROUP, table_op, glue_op, validating, assert_within_ulp + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank, world, local_rank = _rank_env() + os.environ.setdefault("RANK", str(rank)) + os.environ.setdefault("WORLD_SIZE", str(world)) + torch.cuda.set_device(local_rank) + dist.init_process_group( + backend="cuda:nccl,cpu:gloo", + init_method="env://", + rank=rank, + world_size=world, + device_id=torch.device("cuda", local_rank), + ) + RANK, WORLD = rank, world + HOST_GROUP = dist.new_group(backend="gloo") + + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec import ( # noqa: E501 + _draft_exchange, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec.draft_ehproj_table import ( # noqa: E501 + glm_draft_ehproj_table, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec.draft_sample_glue import ( # noqa: E501 + glm_draft_sample_glue, + ) + + table_op, glue_op = glm_draft_ehproj_table, glm_draft_sample_glue + # Spawned, not collected: reach the tree's _validating the plain way. tests/unittest/_torch/modeling_v2 is + # four levels above this directory. + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import assert_within_ulp as _within + from _validating import validating as _validating_cm + + validating, assert_within_ulp = _validating_cm, _within + + hosts: List[Any] = [None] * WORLD + dist.all_gather_object(hosts, socket.gethostname(), group=HOST_GROUP) + fabric = len(set(hosts)) > 1 or os.environ.get("TRTLLM_FORCE_MNNVL_AR", "0") == "1" + assert _draft_exchange.ehproj_table_supports_tp(WORLD), ( + f"the table op does not serve {WORLD} ranks" + ) + max_rows = _draft_exchange.MAX_ROWS + WORKSPACES["table"] = _draft_exchange.DraftExchangeWorkspace( + max_rows * WORLD * _draft_exchange.ehproj_table_slot_bytes(WORLD), + WORLD, + RANK, + local_rank, + fabric, + dist.group.WORLD, + ) + WORKSPACES["glue"] = _draft_exchange.DraftExchangeWorkspace( + max_rows * WORLD * _draft_exchange.sample_glue_slot_bytes(WORLD), + WORLD, + RANK, + local_rank, + fabric, + dist.group.WORLD, + ) + + checks = { + "draft_ehproj_table": (check_table_cells,), + "draft_sample_glue": (check_glue_cells,), + }[entry] + (check_draft_loop_eager_and_graph, check_guards_refuse) + for check in checks: + try: + ran = check() + if check in (check_table_cells, check_glue_cells) and expected_cells >= 0: + assert ran == expected_cells, ( + f"ran {ran} cells, the entry declares {expected_cells}" + ) + except BaseException: + import traceback + + print(f"[rank {RANK}] FAILED {check.__name__}", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + # Exit hard: a rank that leaves an exchange early leaves its peers polling. + os._exit(13) + dist.barrier(group=HOST_GROUP) + torch.cuda.synchronize() + print(f"[rank {RANK}] {entry}: {len(checks)} checks passed", flush=True) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks(entry: str, expected_cells: int) -> None: + """One worker per device in CUDA_VISIBLE_DEVICES; every session is killed on a failure or the deadline.""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the devices this run owns" + devices = [d for d in visible.split(",") if d.strip()] + world_size = len(devices) + assert world_size >= 2, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); an exchange needs 2" + ) + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + ) + # The children rendezvous through torch.distributed's env://, not the launching step's PMI namespace. + env = { + k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_")) + } + print(f"[launcher] spawning {world_size} ranks for {entry}, rendezvous on :{port}", flush=True) + children = [ + subprocess.Popen( + [ + sys.executable, + str(Path(__file__).resolve()), + entry, + str(expected_cells), + _WORKER_FLAG, + ], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + for r in range(world_size) + ] + + def _kill_all() -> None: + for p in children: + try: + os.killpg(p.pid, signal.SIGKILL) + except ProcessLookupError: + pass + for p in children: + p.wait() + + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" + break + for r in sorted(live): + code = children[r].poll() + if code is None: + continue + live.discard(r) + if code != 0: + failure = f"rank {r} exited {code}" + break + time.sleep(0.5) + if failure is not None: + _kill_all() + raise AssertionError(f"the {world_size}-rank {entry} run failed: {failure}") + + +if __name__ == "__main__": + args = [a for a in sys.argv[1:] if a != _WORKER_FLAG] + entry_name = args[0] + cells = int(args[1]) if len(args) > 1 else -1 + if _WORKER_FLAG in sys.argv: + sys.exit(_run_one_rank(entry_name, cells)) + _spawn_ranks(entry_name, cells) + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_rank_job.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_rank_job.py new file mode 100644 index 000000000000..96277c57c6d5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_rank_job.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Run `_draft_step_op_matrix.py` for one entry from pytest, and turn its exit code into an assertion. + +The matrix owns the hard parts -- one rank per device, its own process group, a deadline it enforces by killing +every rank's session -- so this only selects the devices and starts it in a fresh interpreter, which must carry +no CUDA or torch.distributed state of its own (a pytest process that imported tensorrt_llm does). Started by +file path: the ranks reach the target by absolute import, and this tree is not a package. +""" + +from __future__ import annotations + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path + +WORLD_SIZE = 4 +"""Ranks of the single-node run: the 16-rank geometry needs four nodes and an external launcher.""" + +_LAUNCHER_GRACE_S = 300 + + +def _devices() -> str: + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible: + devices = [d for d in visible.split(",") if d.strip()] + else: + import torch + + devices = [str(i) for i in range(torch.cuda.device_count())] + assert len(devices) >= WORLD_SIZE, ( + f"the draft-step exchange runs over {WORLD_SIZE} ranks; " + f"only {len(devices)} device(s) are visible" + ) + return ",".join(devices[:WORLD_SIZE]) + + +def run(entry: str, expected_cells: int) -> None: + """Run the matrix for `entry` over WORLD_SIZE devices, expecting it to drive `expected_cells` cells.""" + launcher = Path(__file__).resolve().with_name("_draft_step_op_matrix.py") + env = dict(os.environ, CUDA_VISIBLE_DEVICES=_devices()) + # The ranks import tensorrt_llm absolutely, and a source checkout is not necessarily installed. + repo_root = Path(__file__).resolve().parents[8] + assert (repo_root / "tensorrt_llm").is_dir(), ( + f"expected the repo root at {repo_root}, found no tensorrt_llm/ there" + ) + env["PYTHONPATH"] = os.pathsep.join( + p for p in (str(repo_root), env.get("PYTHONPATH", "")) if p + ) + spec = importlib.util.spec_from_file_location(launcher.stem, launcher) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + completed = subprocess.run( + [sys.executable, str(launcher), entry, str(expected_cells)], + env=env, + capture_output=True, + text=True, + timeout=module.DEADLINE_S + _LAUNCHER_GRACE_S, + ) + if completed.returncode != 0: + raise AssertionError( + f"{launcher.name} {entry} exited {completed.returncode}\n" + f"--- stdout ---\n{completed.stdout}\n" + f"--- stderr ---\n{completed.stderr}" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_ehproj_table.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_ehproj_table.py new file mode 100644 index 000000000000..b536ec7d560e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_ehproj_table.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the glm_draft_ehproj_table entry's multi-rank certification. + +The cells, a draft loop of the target's shape (eager and as CUDA-graph replays) and the guards run inside +`_draft_step_op_matrix.py` beside this file, one process per rank over a 4-rank group; see that file for the +16-rank run over four nodes. +""" + +import _draft_step_rank_job +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec.draft_ehproj_table import ( # noqa: E501 + glm_draft_ehproj_table as op, +) +from tensorrt_llm._utils import get_sm_version + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" +) + + +# Each run spawns its own 4 rank processes over the visible devices: under xdist several workers would fight +# for the same GPUs. +@pytest.mark.no_xdist +def test_draft_ehproj_table_op_matrix() -> None: + assert op.CELLS, "the entry declares no cells" + _draft_step_rank_job.run("draft_ehproj_table", expected_cells=len(op.CELLS)) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_sample_glue.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_sample_glue.py new file mode 100644 index 000000000000..faf00183c753 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_sample_glue.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the glm_draft_sample_glue entry's multi-rank certification. + +The cells, a draft loop of the target's shape (eager and as CUDA-graph replays) and the guards run inside +`_draft_step_op_matrix.py` beside this file, one process per rank over a 4-rank group; see that file for the +16-rank run over four nodes. +""" + +import _draft_step_rank_job +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec.draft_sample_glue import ( # noqa: E501 + glm_draft_sample_glue as op, +) +from tensorrt_llm._utils import get_sm_version + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" +) + + +# Each run spawns its own 4 rank processes over the visible devices: under xdist several workers would fight +# for the same GPUs. +@pytest.mark.no_xdist +def test_draft_sample_glue_op_matrix() -> None: + assert op.CELLS, "the entry declares no cells" + _draft_step_rank_job.run("draft_sample_glue", expected_cells=len(op.CELLS)) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_token_table.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_token_table.py new file mode 100644 index 000000000000..a92b363d627b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_token_table.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the construction-time operands of the glm_draft_ehproj_table entry.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec._draft_token_table import ( # noqa: E501 + build_draft_token_table, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" +) + +HIDDEN = 6144 +EPS = 1e-5 + + +@pytest.mark.parametrize("tp_size,tp_rank", [(16, 0), (16, 11), (4, 3)]) +def test_table_is_the_token_half_of_eh_proj(tp_size, tp_rank) -> None: + """Each table row is eh_proj's embedding half of one token of the rank's shard. + + Against an FP32 RMSNorm and product, rounded where the builder rounds (the normalized row and the product to + BF16), the rows differ only where the order of a 6144-term FP32 sum or the kernel's approximate reciprocal + square root moves a BF16 rounding: one ulp of that element. So 1 ulp against the row scale and 0.25 ulp of + relative RMS. The hidden-half rows are a slice of the weight, compared exactly. + """ + g = torch.Generator(device="cuda").manual_seed(tp_size * 100 + tp_rank) + vocab_shard = 96 + weight = ( + torch.randn(HIDDEN, 2 * HIDDEN, generator=g, device="cuda") * HIDDEN**-0.5 + ).to(torch.bfloat16) + embed = torch.randn(vocab_shard * tp_size, HIDDEN, generator=g, device="cuda").to( + torch.bfloat16 + ) + enorm = (1.0 + 0.1 * torch.randn(HIDDEN, generator=g, device="cuda")).to(torch.bfloat16) + + table, hidden_rows = build_draft_token_table( + weight, embed, enorm, EPS, tp_rank, tp_size, vocab_shard + ) + + rows = embed[tp_rank * vocab_shard : (tp_rank + 1) * vocab_shard].float() + rstd = torch.rsqrt(rows.pow(2).mean(dim=-1, keepdim=True) + EPS) + normed = ((rows * rstd) * enorm.float()).to(torch.bfloat16).float() + with true_fp32_matmul(): + expected = (normed @ weight[:, :HIDDEN].float().t()).to(torch.bfloat16) + assert table.shape == (vocab_shard, HIDDEN) and table.dtype == torch.bfloat16 + assert_within_ulp(table, expected, element_ulp=1.0, rms_ulp=0.25) + share = HIDDEN // tp_size + torch.testing.assert_close( + hidden_rows, weight[tp_rank * share : (tp_rank + 1) * share, HIDDEN:], rtol=0, atol=0 + ) + assert hidden_rows.is_contiguous() From 4be90c52360b0f7b74a8117646517de3ee4911ae Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:07:18 -0700 Subject: [PATCH 27/53] Add the glm_moe_dsa tp16 target's MNNVL one-shot hops trtllm::glm_mnnvl_hop runs a BF16 one-shot MNNVL all-reduce with the residual add and RMSNorm on the engine's MNNVL all-reduce workspace: - every rank stores its rows through the multicast address and polls all slots in one loop; the sum is FP32 in rank order, from zero per chunk of 8 ranks; - the residual and gamma are copied into shared memory before the exchange, and the warp sums of squares travel to every CTA of the cluster through a transaction barrier (st.async); - each CTA counts its own arrival, and the thread that rotates the Lamport flags reads the arrival counter while it polls and rotates them right after its poll; - the PDL dependents launch right after the grid dependency wait; - with shard_published the producer of the input (the o_proj Lamport publish) has already stored this rank's rows into the workspace: the kernel polls every slot, reads the flags and the residual with volatile loads before any wait, and waits for its primary at its end. trtllm::glm_mnnvl_moe_finalize_hop fuses the TRTLLM-Gen MoE finalize (top-8, the shared experts' output as addend) into the same hop, and trtllm::glm_mnnvl_hop_supported states the shapes both serve: 2, 4, 8 or 16 ranks, one-shot (1 MiB per rank, 1.5 MiB on 16 ranks), one row per cluster of 2 to 8 CTAs of at most 128 threads covering the row. GlmLamportFlags reads and rotates the flags in the layout of common::LamportFlags, with per-CTA arrivals, so the hops and the built-in MNNVL all-reduce follow each other on one workspace. The entries mnnvl_hop and mnnvl_moe_finalize_hop carry the c1 cells (hidden 6144; 8 rows, and 1 row on draft steps). _mnnvl_workspace returns the hop operands and the o_proj publish target from the engine's workspace. The 4-rank op matrices check the cells against FP32 references, against the built-in one-shot all-reduce bit for bit, the fused finalize against finalize + hop bit for bit, the flags, a published shard, CUDA graph replays, late ranks and the guards. Signed-off-by: Dongfeng Yu --- .../kernels/glmMoeDsa/glmLamport.cuh | 183 ++++ .../kernels/glmMoeDsa/glmMnnvlHop.cu | 799 ++++++++++++++++++ .../kernels/glmMoeDsa/glmMnnvlHop.h | 99 +++ cpp/tensorrt_llm/thop/CMakeLists.txt | 1 + cpp/tensorrt_llm/thop/glmMnnvlHopOp.cpp | 217 +++++ .../ops/comm/_mnnvl_workspace.py | 86 ++ .../ops/comm/mnnvl_hop.py | 135 +++ .../ops/comm/mnnvl_moe_finalize_hop.py | 140 +++ .../comm/_mnnvl_hop_op_matrix.py | 420 +++++++++ .../comm/_mnnvl_moe_finalize_hop_op_matrix.py | 354 ++++++++ .../comm/_mnnvl_ranks.py | 221 +++++ .../comm/_rank_job.py | 75 ++ .../test_modeling_v2_mnnvl_hop_op_matrix.py | 24 + ...ing_v2_mnnvl_moe_finalize_hop_op_matrix.py | 24 + 14 files changed, 2778 insertions(+) create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmLamport.cuh create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.cu create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.h create mode 100644 cpp/tensorrt_llm/thop/glmMnnvlHopOp.cpp create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/_mnnvl_workspace.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_hop.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_moe_finalize_hop.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_hop_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_moe_finalize_hop_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_ranks.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_rank_job.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_hop_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_moe_finalize_hop_op_matrix.py diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLamport.cuh b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLamport.cuh new file mode 100644 index 000000000000..c7461c9dbc00 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLamport.cuh @@ -0,0 +1,183 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/common/lamportUtils.cuh" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +//! \brief The Lamport flags of an MNNVL workspace, read and rotated as common::LamportFlags does. +//! +//! The flags are nine 32-bit words: [0] the current buffer, [1] the dirty buffer (the previous call's), [2] bytes per +//! buffer, [3] stages of the dirty buffer, [4..7] bytes to clear per stage of the dirty buffer, [8] the arrival +//! counter. A kernel publishes into the current buffer, empties the dirty one, and the thread that rotates the flags +//! makes the current buffer the dirty one once every CTA has counted its arrival. Each CTA counts its arrival with a +//! CTA barrier and a release reduction, as common::LamportFlags does, so kernels of either type follow each other on +//! one workspace. On top of common::LamportFlags: a volatile read of the flags, the arrival counter, the rotation +//! without its wait, and the fill of the dirty buffer with a value other than -0.0. +//! +//! Like common::LamportFlags, the clear and the fill need every thread of a 1D-block grid. +template +class GlmLamportFlags +{ +public: + //! \param bufferFlags The nine flag words of the workspace. + //! \param numStages Stages this kernel uses in the current buffer. + //! \param volatileRead Read the flags with volatile loads, which do not go through L1: for a kernel that reads them + //! before its grid dependency wait. + __device__ GlmLamportFlags(uint32_t* bufferFlags, uint32_t numStages, bool volatileRead) + : mBufferFlagsPtr(bufferFlags) + , mFlagAccessPtr(&bufferFlags[kArrivalCounterWord]) + { + uint4 flag; + uint4 clear; + if (volatileRead) + { + auto const flagWords = common::loadPackedVolatile(bufferFlags); + auto const clearWords = common::loadPackedVolatile(bufferFlags + 4); + flag = make_uint4(flagWords.words[0], flagWords.words[1], flagWords.words[2], flagWords.words[3]); + clear = make_uint4(clearWords.words[0], clearWords.words[1], clearWords.words[2], clearWords.words[3]); + } + else + { + flag = reinterpret_cast(bufferFlags)[0]; + clear = reinterpret_cast(bufferFlags)[1]; + } + mCurrentIndex = flag.x; + mDirtyIndex = flag.y; + mCurBufferLayout.numStages = numStages; + mCurBufferLayout.bytesPerBuffer = flag.z; + mDirtyBufferLayout.bytesPerBuffer = flag.z; + mDirtyBufferLayout.numStages = flag.w; + mBytesToClear[0] = clear.x; + mBytesToClear[1] = clear.y; + mBytesToClear[2] = clear.z; + mBytesToClear[3] = clear.w; + } + + //! The base of stage stageIdx of the current buffer. + [[nodiscard]] __device__ void* getCurLamportBuf(void* bufferBasePtr, int stageIdx = 0) const + { + return mCurBufferLayout.getStagePtr(bufferBasePtr, mCurrentIndex, stageIdx); + } + + //! Writes the Lamport empty value (-0.0) into every stage of the dirty buffer. + __device__ void clearDirtyLamportBuf(void* bufferBasePtr) + { + fillDirtyLamportBuf(bufferBasePtr, common::getPackedLamportInit()); + } + + //! Writes fillValue into every stage of the dirty buffer, for a workspace whose empty value is not -0.0. + __device__ void fillDirtyLamportBuf(void* bufferBasePtr, PackedType fillValue) + { + uint32_t const globalCtaIdx = blockIdx.x * gridDim.y + blockIdx.y; + uint32_t const globalTid = globalCtaIdx * blockDim.x + threadIdx.x; + uint32_t const numThreads = gridDim.x * gridDim.y * blockDim.x; + for (uint32_t stage = 0; stage < mDirtyBufferLayout.numStages; stage++) + { + uint32_t const fillBoundary = (mBytesToClear[stage] + sizeof(PackedType) - 1) / sizeof(PackedType); + auto* dirty + = static_cast(mDirtyBufferLayout.getStagePtr(bufferBasePtr, mDirtyIndex, stage)); + for (uint32_t packedIdx = globalTid; packedIdx < fillBoundary; packedIdx += numThreads) + { + dirty[packedIdx] = fillValue; + } + } + } + + //! Counts this CTA's arrival: a barrier of the CTA's threads (named barrier 1), then a release reduction on the + //! arrival counter. Every flag read of the CTA precedes it. For kernels of SM 90 and newer. + __device__ void ctaArrive() + { + uint32_t const barrierThreads + = ((static_cast(blockDim.x) + common::kWARP_SIZE - 1U) / common::kWARP_SIZE) * common::kWARP_SIZE; + if (threadIdx.x < common::kWARP_SIZE) + { + asm volatile("barrier.cta.sync 1, %0;" ::"r"(barrierThreads) : "memory"); + if (threadIdx.x == 0) + { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) + asm volatile("red.async.release.global.gpu.add.u32 [%0], %1;" ::"l"(mFlagAccessPtr), "r"(1) : "memory"); +#else + asm volatile("red.release.global.gpu.add.u32 [%0], %1;" ::"l"(mFlagAccessPtr), "r"(1) : "memory"); +#endif + } + } + else + { + asm volatile("barrier.cta.arrive 1, %0;" ::"r"(barrierThreads) : "memory"); + } + } + + //! The arrival counter that ctaArrive increments. + [[nodiscard]] __device__ uint32_t const* arrivalCounter() const + { + return mFlagAccessPtr; + } + + //! The rotation of waitAndUpdate without its wait: the current buffer becomes the dirty one. Only for the thread + //! that waitAndUpdate elects, once it has observed every CTA's arrival with acquire semantics. + __device__ void update(uint4 bytesToClearPerStage) + { + uint4* flagPtr = reinterpret_cast(mBufferFlagsPtr); + flagPtr[0] = {(mCurrentIndex + 1) % common::LamportBufferLayout::sNumLamportBuffers, mCurrentIndex, + mCurBufferLayout.bytesPerBuffer, mCurBufferLayout.numStages}; + flagPtr[1] = bytesToClearPerStage; + *mFlagAccessPtr = 0; + } + + //! Thread 0 of CTA 0 waits until every CTA has arrived, then rotates the flags. + __device__ void waitAndUpdate(uint4 bytesToClearPerStage) + { + if (threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0) + { + uint32_t const targetCount = gridDim.x * gridDim.y * gridDim.z; + uint32_t arrivedCount = 0; + do + { + asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" + : "=r"(arrivedCount) + : "l"(mFlagAccessPtr) + : "memory"); + } while (arrivedCount < targetCount); + update(bytesToClearPerStage); + } + } + +private: + static constexpr int kArrivalCounterWord = 8; + + uint32_t* mBufferFlagsPtr; + uint32_t* mFlagAccessPtr; + uint32_t mCurrentIndex; + uint32_t mDirtyIndex; + uint32_t mBytesToClear[4]; + common::LamportBufferLayout mCurBufferLayout; + common::LamportBufferLayout mDirtyBufferLayout; +}; + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.cu new file mode 100644 index 000000000000..41e9f7009465 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.cu @@ -0,0 +1,799 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaTypeUtils.cuh" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/common/lamportUtils.cuh" +#include "tensorrt_llm/common/reduceKernelUtils.cuh" +#include "tensorrt_llm/kernels/glmMoeDsa/glmLamport.cuh" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ +namespace +{ + +using common::cuda_cast; + +// Eight BF16 elements, one 16-byte access. +constexpr int kEltsPerThread = static_cast(sizeof(float4) / sizeof(__nv_bfloat16)); +constexpr int kWarpSize = 32; +constexpr int kLog2WarpSize = 5; +constexpr int kLaneIdMask = kWarpSize - 1; +// The fast residual + RMSNorm epilogue takes CTAs of up to four whole warps in clusters of 2 to 8 CTAs. +constexpr int kMaxThreads = 128; +constexpr int kMaxWarps = kMaxThreads / kWarpSize; +constexpr int kMaxClusterSize = 8; +// Ranks summed from zero before their sum joins the total. +constexpr int kRankChunk = 8; +// Lamport data per rank up to which an all-reduce runs one-shot: 1 MiB (the MNNVL sweep's threshold), and on 16 ranks +// 1.5 MiB, which keeps an 8-row hidden-6144 BF16 step one-shot. +constexpr int64_t kOneShotBytes = 64 * 1024 * 8 * 2; +constexpr int64_t kOneShotBytes16Ranks = 3 * 512 * 1024; + +union PackedBf16x8 +{ + float4 packed; + __nv_bfloat16 elements[kEltsPerThread]; +}; + +struct HopKernelParams +{ + __nv_bfloat16* normOut; + __nv_bfloat16* residualOut; + __nv_bfloat16 const* shard; + __nv_bfloat16 const* residualIn; + __nv_bfloat16 const* gamma; + __nv_bfloat16** bufferPtrs; + __nv_bfloat16* mcastPtr; + uint32_t* bufferFlags; + int numTokens; + int tokenDim; + int rank; + float epsilon; +}; + +struct MoeFinalizeKernelParams +{ + __nv_bfloat16 const* fc2Output; + __nv_bfloat16 const* expertWeights; + int32_t const* expandedIdxToPermutedIdx; + __nv_bfloat16 const* addend; + int hiddenDimPadded; +}; + +// Shared memory of the fast epilogue: each thread's residual and gamma vectors, copied before the Lamport exchange; the +// warp sums of squares of every CTA of the cluster, kMaxWarps per CTA; and the transaction barrier that completes once +// all of them have landed. +struct FastEpilogueStorage +{ + float4 residual[kMaxThreads]; + float4 gamma[kMaxThreads]; + float4 warpSums[kMaxClusterSize]; + uint64_t warpSumsBarrier; +}; + +// Only the kernels that call this function allocate the storage. +__device__ __forceinline__ FastEpilogueStorage& fastEpilogueStorage() +{ + __shared__ FastEpilogueStorage storage; + return storage; +} + +__device__ __forceinline__ void copyF4(float4* dstShared, __nv_bfloat16 const* src) +{ + __pipeline_memcpy_async(dstShared, reinterpret_cast(src), sizeof(float4)); +} + +// -0.0 is the empty value of the Lamport buffers: a payload element of -0.0 is sent as +0.0. +__device__ __forceinline__ void sanitizeLamportPayload(PackedBf16x8& value) +{ +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + if (common::isNegZero(value.elements[i])) + { + value.elements[i] = cuda_cast<__nv_bfloat16, float>(0.F); + } + } +} + +__device__ __forceinline__ void accumulatePacked(float (&accum)[kEltsPerThread], PackedBf16x8 const& value) +{ +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + accum[i] += cuda_cast(value.elements[i]); + } +} + +__device__ __forceinline__ uint32_t loadArrivalsRelaxed(uint32_t const* arrivalCounter) +{ + uint32_t arrivals; + asm volatile("ld.relaxed.gpu.global.u32 %0, [%1];" : "=r"(arrivals) : "l"(arrivalCounter) : "memory"); + return arrivals; +} + +// Strong loads at GPU scope (L2, never a line that L1 kept from an earlier use), for reads before the grid dependency +// wait. +__device__ __forceinline__ int32_t loadRelaxedGpu(int32_t const* ptr) +{ + int32_t value; + asm volatile("ld.relaxed.gpu.global.b32 %0, [%1];" : "=r"(value) : "l"(ptr)); + return value; +} + +// The bits of a BF16 value. +__device__ __forceinline__ uint16_t loadRelaxedGpu(__nv_bfloat16 const* ptr) +{ + uint16_t bits; + asm volatile("ld.relaxed.gpu.global.b16 %0, [%1];" : "=h"(bits) : "l"(ptr)); + return bits; +} + +// The one-shot exchange in one poll loop: every iteration loads the slots of all WorldSize ranks together until none +// holds the Lamport empty value. With LocalValue the slot of localRank is not waited for: localValue, the bits this +// kernel stored there, stands in for it. With TrackArrivals a thread with a non-null arrivalCounter reads the counter +// (relaxed) at the start of every iteration, so that the last value it keeps in *arrivals is about as recent as the +// peers' data. The sum: the ranks of each chunk of kRankChunk summed in rank order from zero in FP32, the chunk sums +// added to zero in order, rounded to BF16. The WorldSize values in flight need the registers of a launch bound of +// kMaxThreads. +template +__device__ __forceinline__ PackedBf16x8 reduceOneshotPollAll(PackedBf16x8 const& localValue, int localRank, + __nv_bfloat16 const* stagePtrLocal, int token, int tokenDim, int packedIdx, uint32_t const* arrivalCounter, + uint32_t* arrivals) +{ + constexpr int kChunk = WorldSize < kRankChunk ? WorldSize : kRankChunk; + static_assert(WorldSize % kChunk == 0, "The ranks must split into whole chunks"); + __nv_bfloat16 const* slots = &stagePtrLocal[token * tokenDim * WorldSize + packedIdx * kEltsPerThread]; + PackedBf16x8 values[WorldSize]; + while (true) + { + if constexpr (TrackArrivals) + { + if (arrivalCounter != nullptr) + { + *arrivals = loadArrivalsRelaxed(arrivalCounter); + } + } + bool valid = true; +#pragma unroll + for (int r = 0; r < WorldSize; r++) + { + auto const loaded = common::loadPackedVolatile(&slots[r * tokenDim]); + values[r].packed = loaded.packed; + valid &= (LocalValue && r == localRank) | !common::isLamportDirty(loaded); + } + if (valid) + { + break; + } + } + float accum[kEltsPerThread]; +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + accum[i] = 0.F; + } +#pragma unroll + for (int rankBase = 0; rankBase < WorldSize; rankBase += kChunk) + { + float chunkAccum[kEltsPerThread]; +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + chunkAccum[i] = 0.F; + } +#pragma unroll + for (int rr = 0; rr < kChunk; rr++) + { + int const r = rankBase + rr; + accumulatePacked(chunkAccum, LocalValue && r == localRank ? localValue : values[r]); + } +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + accum[i] += chunkAccum[i]; + } + } + PackedBf16x8 packedAccum; +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + packedAccum.elements[i] = cuda_cast<__nv_bfloat16, float>(accum[i]); + } + return packedAccum; +} + +// Called by one thread per CTA before every thread's relaxed arrive on the cluster barrier: the barrier expects one +// arrival (this one) and the bytes of the warp sums of every CTA of the cluster. That cluster barrier phase publishes +// the initialization; every thread waits for the phase before it sends. +__device__ __forceinline__ void initFastEpilogueBarrier() +{ + namespace cg = cooperative_groups; + uint32_t const bytes + = static_cast(cg::this_cluster().num_blocks() * (blockDim.x >> kLog2WarpSize) * sizeof(float)); + uint32_t const barrier = static_cast(__cvta_generic_to_shared(&fastEpilogueStorage().warpSumsBarrier)); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(barrier) : "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(barrier), "r"(bytes) : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); +} + +// Sum of threadSum over all threads of the cluster. Each warp reduces its sums with the warpReduceSum butterfly (every +// lane ends with the same bits); every warp sends its sum to every CTA of the cluster with st.async, and each CTA waits +// on its transaction barrier until all warp sums have landed. Per CTA the warp sums are added as (w0 + w2) + (w1 + w3) +// without the absent terms, and the CTA sums in block-rank order from zero: the order of blockReduceSum +// followed by the cluster loop of the one-shot all-reduce's default epilogue. Needs whole warps, at most kMaxWarps +// warps per CTA, 2 to kMaxClusterSize CTAs per cluster, the barrier initialized by initFastEpilogueBarrier and the +// cluster barrier phase that published it completed. +__device__ __forceinline__ float clusterSumOfSquaresFast(float threadSum) +{ + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + auto& storage = fastEpilogueStorage(); + int const numBlocks = static_cast(cluster.num_blocks()); + int const blockRank = static_cast(cluster.block_rank()); + int const numWarps = static_cast(blockDim.x >> kLog2WarpSize); + int const warpId = static_cast(threadIdx.x >> kLog2WarpSize); + int const laneId = static_cast(threadIdx.x & kLaneIdMask); + + float const warpSum = common::warpReduceSum(threadSum); + uint32_t const barrier = static_cast(__cvta_generic_to_shared(&storage.warpSumsBarrier)); + if (laneId < numBlocks) + { + // Lane l sends to CTA l. + uint32_t const slot = static_cast(__cvta_generic_to_shared( + reinterpret_cast(storage.warpSums) + blockRank * kMaxWarps + warpId)); + uint32_t peerSlot; + uint32_t peerBarrier; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(peerSlot) : "r"(slot), "r"(laneId)); + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(peerBarrier) : "r"(barrier), "r"(laneId)); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(peerSlot), + "r"(__float_as_uint(warpSum)), "r"(peerBarrier) + : "memory"); + } + uint32_t complete = 0; + while (!complete) + { + asm volatile( + "{\n" + ".reg .pred done;\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 done, [%1], 0;\n" + "selp.u32 %0, 1, 0, done;\n" + "}\n" + : "=r"(complete) + : "r"(barrier) + : "memory"); + } + + float fullSum = 0.F; +#pragma unroll + for (int block = 0; block < kMaxClusterSize; ++block) + { + if (block < numBlocks) + { + float4 const sums = storage.warpSums[block]; + float const even = numWarps > 2 ? sums.x + sums.z : sums.x; + float const odd = numWarps > 3 ? sums.y + sums.w : sums.y; + fullSum += numWarps > 1 ? even + odd : even; + } + } + return fullSum; +} + +// The rotation of the Lamport flags by the thread that GlmLamportFlags::waitAndUpdate elects. arrivals is the last +// value that thread read from the arrival counter with a relaxed load. Once it covers every CTA, the fence turns that +// read into an acquire of every CTA's release arrival and the flags rotate at once; otherwise the thread waits on the +// counter as waitAndUpdate does. +__device__ __forceinline__ void rotateLamportFlags( + GlmLamportFlags& flag, uint32_t arrivals, uint4 bytesToClearPerStage) +{ + uint32_t const numCtas = gridDim.x * gridDim.y * gridDim.z; + if (arrivals >= numCtas) + { + asm volatile("fence.acquire.gpu;" ::: "memory"); + } + else + { + do + { + asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" + : "=r"(arrivals) + : "l"(flag.arrivalCounter()) + : "memory"); + } while (arrivals < numCtas); + } + flag.update(bytesToClearPerStage); +} + +// After the poll of the rotating thread: rotate now if the last count it read covers every CTA; otherwise read the +// counter once more, while the epilogue runs, and leave the rotation to the end of the kernel. +__device__ __forceinline__ void rotateIfAllArrived( + GlmLamportFlags& flag, bool& rotationPending, uint32_t& arrivals, uint4 bytesToClearPerStage) +{ + if (rotationPending) + { + if (arrivals >= gridDim.x * gridDim.y * gridDim.z) + { + rotateLamportFlags(flag, arrivals, bytesToClearPerStage); + rotationPending = false; + } + else + { + arrivals = loadArrivalsRelaxed(flag.arrivalCounter()); + } + } +} + +// The fast residual + RMSNorm epilogue of both hops, after the copies of the residual and gamma have completed: the +// residual added in BF16, the BF16 squares summed in FP32 in element order and over the cluster, (x * rsqrt) * gamma +// rounded to BF16, the residual output stored after the exchange of the sums (the complete-tx of each st.async +// releases every earlier store of the thread), and the rotation of the flags if it is still pending. +__device__ __forceinline__ void hopEpilogue(HopKernelParams const& params, GlmLamportFlags& flag, + PackedBf16x8 packedAccum, PackedBf16x8 const& residualIn, int threadOffset, bool rotationPending, uint32_t arrivals, + uint4 bytesToClearPerStage) +{ +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + packedAccum.elements[i] += residualIn.elements[i]; + } + PackedBf16x8 gamma; + gamma.packed = fastEpilogueStorage().gamma[threadIdx.x]; + float threadSum = 0.F; +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + threadSum += cuda_cast(packedAccum.elements[i] * packedAccum.elements[i]); + } + // Every thread arrived on the cluster barrier before its grid dependency wait; waiting for that phase publishes + // every CTA's transaction barrier before any thread sends. + __cluster_barrier_wait(); + float const fullSum = clusterSumOfSquaresFast(threadSum); + reinterpret_cast(¶ms.residualOut[threadOffset])[0] = packedAccum.packed; + float const rcpRms = rsqrtf(fullSum / params.tokenDim + params.epsilon); +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + packedAccum.elements[i] = cuda_cast<__nv_bfloat16, float>( + cuda_cast(packedAccum.elements[i]) * rcpRms + * cuda_cast(gamma.elements[i])); + } + reinterpret_cast(¶ms.normOut[threadOffset])[0] = packedAccum.packed; + if (rotationPending) + { + rotateLamportFlags(flag, arrivals, bytesToClearPerStage); + } +} + +// One row per cluster, eight elements per thread, the clusters' CTAs covering the row exactly. Each CTA counts its +// own arrival (GlmLamportFlags) and the thread that rotates the flags tracks the arrival counter in its poll. The +// transaction barrier is published with a relaxed cluster arrive before the grid dependency wait. The PDL dependents +// launch right after the grid dependency wait (at the start with a published shard) and must read the outputs only +// after their own wait. +// ShardPublished: the producer of the input has stored this rank's rows into the current Lamport buffer of every rank, +// sanitized, where this kernel would store them, so the kernel neither loads nor stores them and polls every rank's +// slot, its own included. It does not wait for its primary before the exchange: the producer releases it only after its +// own grid dependency wait, so every grid before the producer has completed, the previous hop (the flags, the residual) +// included. The flags and the residual are read with volatile loads, gamma (a weight) with the epilogue's copy. A CTA +// counts its arrival only once all of its slots are in: the rotation after the last arrival must not precede a producer +// CTA's read of the flags. The grid waits for its primary at its end, so that its completion still implies the +// producer's. +template +__global__ void __launch_bounds__(kMaxThreads) mnnvlHopKernel(HopKernelParams params) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + int const packedIdx = static_cast(cluster.thread_rank()); + int const token = static_cast(blockIdx.x); + int const threadOffset = token * params.tokenDim + packedIdx * kEltsPerThread; + + // The transaction barrier is this CTA's shared memory, which the previous grid does not touch: initialize it and + // publish it to the cluster while the grid dependency wait is pending. + if (threadIdx.x == 0) + { + initFastEpilogueBarrier(); + } + asm volatile("barrier.cluster.arrive.relaxed;" ::: "memory"); + if constexpr (!ShardPublished) + { + cudaGridDependencySynchronize(); + } + cudaTriggerProgrammaticLaunchCompletion(); + + GlmLamportFlags flag(params.bufferFlags, /*numStages=*/1, /*volatileRead=*/ShardPublished); + auto* stagePtrMcast = static_cast<__nv_bfloat16*>(flag.getCurLamportBuf(params.mcastPtr)); + auto* stagePtrLocal = static_cast<__nv_bfloat16*>(flag.getCurLamportBuf(params.bufferPtrs[params.rank])); + auto& storage = fastEpilogueStorage(); + + // ==================== Broadcast this rank's rows ==================== + PackedBf16x8 val{}; + [[maybe_unused]] PackedBf16x8 residualPublished; + if constexpr (ShardPublished) + { + residualPublished.packed = common::loadPackedVolatile(¶ms.residualIn[threadOffset]).packed; + copyF4(&storage.gamma[threadIdx.x], ¶ms.gamma[packedIdx * kEltsPerThread]); + } + else + { + val.packed = *reinterpret_cast(¶ms.shard[threadOffset]); + // The residual and gamma do not depend on the exchange: copy them now, so that their latency overlaps the + // broadcast and the poll instead of following the reduction. + copyF4(&storage.residual[threadIdx.x], ¶ms.residualIn[threadOffset]); + copyF4(&storage.gamma[threadIdx.x], ¶ms.gamma[packedIdx * kEltsPerThread]); + sanitizeLamportPayload(val); + reinterpret_cast( + &stagePtrMcast[token * params.tokenDim * WorldSize + params.rank * params.tokenDim])[packedIdx] + = val.packed; + } + __pipeline_commit(); + if constexpr (!ShardPublished) + { + flag.ctaArrive(); + } + // ==================== Empty the previous call's buffer ==================== + flag.clearDirtyLamportBuf(params.bufferPtrs[params.rank]); + + // ==================== Reduction ==================== + // The thread that waitAndUpdate would elect tracks the arrival counter and rotates the flags. + bool rotationPending = threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0; + uint32_t arrivals = 0; + PackedBf16x8 packedAccum; + if constexpr (ShardPublished) + { + packedAccum = reduceOneshotPollAll( + val, params.rank, stagePtrLocal, token, params.tokenDim, packedIdx, nullptr, nullptr); + flag.ctaArrive(); + } + else + { + packedAccum = reduceOneshotPollAll(val, params.rank, stagePtrLocal, token, + params.tokenDim, packedIdx, rotationPending ? flag.arrivalCounter() : nullptr, &arrivals); + } + uint4 const bytesToClear{ + static_cast(params.numTokens * params.tokenDim * WorldSize * sizeof(__nv_bfloat16)), 0, 0, 0}; + rotateIfAllArrived(flag, rotationPending, arrivals, bytesToClear); + + // ==================== Residual + RMSNorm ==================== + // Each thread reads only the copies it issued, so their completion is the only wait needed. + __pipeline_wait_prior(0); + PackedBf16x8 residualIn; + if constexpr (ShardPublished) + { + residualIn = residualPublished; + } + else + { + residualIn.packed = storage.residual[threadIdx.x]; + } + hopEpilogue(params, flag, packedAccum, residualIn, threadOffset, rotationPending, arrivals, bytesToClear); + if constexpr (ShardPublished) + { + // The completion of this grid implies its primary's, as when it waits at its start. + cudaGridDependencySynchronize(); + } +#endif +} + +// The routed experts' finalize step fused into the hop (BF16, top-8): each thread computes its eight elements of the +// finalize output as the TRTLLM-Gen finalize does (the token's eight FC2 rows weighted in k order with an FMA, -1 slots +// skipped, rounded to BF16; the addend added in FP32 and the sum rounded once more) and publishes them as its rows. +// Everything else is mnnvlHopKernel without a published shard. Eight threads per CTA read the token's FC2 row indices +// and routing weights before the grid dependency wait: the MoE's routing wrote them, and it has completed when this +// grid launches because FC1 and FC2 trigger their dependents only after their own waits. The launch bound lets the +// eight FC2 rows of a thread be in flight in registers. +template +__global__ void __launch_bounds__(kMaxThreads) mnnvlMoeFinalizeHopKernel( + HopKernelParams params, MoeFinalizeKernelParams moe) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + int const packedIdx = static_cast(cluster.thread_rank()); + int const token = static_cast(blockIdx.x); + int const threadOffset = token * params.tokenDim + packedIdx * kEltsPerThread; + __shared__ int32_t sPermutedIdx[kMnnvlHopMoeTopK]; + __shared__ uint16_t sExpertWeightBits[kMnnvlHopMoeTopK]; + + // The transaction barrier is this CTA's shared memory: initialize it and publish it to the cluster while the grid + // dependency wait is pending. + if (threadIdx.x == 0) + { + initFastEpilogueBarrier(); + } + asm volatile("barrier.cluster.arrive.relaxed;" ::: "memory"); + bool const isLoader = threadIdx.x < kMnnvlHopMoeTopK; + int32_t permutedIdxPart = -1; + uint16_t expertWeightBitsPart = 0; + if (isLoader) + { + int const expandedIdx = token * kMnnvlHopMoeTopK + static_cast(threadIdx.x); + permutedIdxPart = loadRelaxedGpu(&moe.expandedIdxToPermutedIdx[expandedIdx]); + expertWeightBitsPart = loadRelaxedGpu(&moe.expertWeights[expandedIdx]); + } + cudaGridDependencySynchronize(); + cudaTriggerProgrammaticLaunchCompletion(); + + GlmLamportFlags flag(params.bufferFlags, /*numStages=*/1, /*volatileRead=*/false); + auto* stagePtrMcast = static_cast<__nv_bfloat16*>(flag.getCurLamportBuf(params.mcastPtr)); + auto* stagePtrLocal = static_cast<__nv_bfloat16*>(flag.getCurLamportBuf(params.bufferPtrs[params.rank])); + auto& storage = fastEpilogueStorage(); + copyF4(&storage.residual[threadIdx.x], ¶ms.residualIn[threadOffset]); + copyF4(&storage.gamma[threadIdx.x], ¶ms.gamma[packedIdx * kEltsPerThread]); + __pipeline_commit(); + if (isLoader) + { + sPermutedIdx[threadIdx.x] = permutedIdxPart; + sExpertWeightBits[threadIdx.x] = expertWeightBitsPart; + } + __syncthreads(); + + // ==================== Finalize: this thread's elements of the routed sum plus the addend ==================== + int permutedIdx[kMnnvlHopMoeTopK]; +#pragma unroll + for (int k = 0; k < kMnnvlHopMoeTopK; k++) + { + permutedIdx[k] = sPermutedIdx[k]; + } + PackedBf16x8 addend; + addend.packed = *reinterpret_cast(&moe.addend[threadOffset]); + PackedBf16x8 rows[kMnnvlHopMoeTopK]; +#pragma unroll + for (int k = 0; k < kMnnvlHopMoeTopK; k++) + { + int64_t const rowOffset + = static_cast(permutedIdx[k]) * moe.hiddenDimPadded + packedIdx * kEltsPerThread; + rows[k].packed = permutedIdx[k] == -1 ? make_float4(0.F, 0.F, 0.F, 0.F) + : *reinterpret_cast(&moe.fc2Output[rowOffset]); + } + float routed[kEltsPerThread]; +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + routed[i] = 0.F; + } +#pragma unroll + for (int k = 0; k < kMnnvlHopMoeTopK; k++) + { + if (permutedIdx[k] == -1) + { + continue; + } + float const weight = cuda_cast(__ushort_as_bfloat16(sExpertWeightBits[k])); +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + routed[i] = __fmaf_rn(weight, cuda_cast(rows[k].elements[i]), routed[i]); + } + } + PackedBf16x8 val; +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + val.elements[i] = cuda_cast<__nv_bfloat16, float>(cuda_cast(addend.elements[i]) + + cuda_cast(cuda_cast<__nv_bfloat16, float>(routed[i]))); + } + + // ==================== Broadcast this rank's rows ==================== + sanitizeLamportPayload(val); + reinterpret_cast( + &stagePtrMcast[token * params.tokenDim * WorldSize + params.rank * params.tokenDim])[packedIdx] + = val.packed; + flag.ctaArrive(); + flag.clearDirtyLamportBuf(params.bufferPtrs[params.rank]); + + // ==================== Reduction ==================== + bool rotationPending = threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0; + uint32_t arrivals = 0; + PackedBf16x8 const packedAccum = reduceOneshotPollAll(val, params.rank, stagePtrLocal, + token, params.tokenDim, packedIdx, rotationPending ? flag.arrivalCounter() : nullptr, &arrivals); + uint4 const bytesToClear{ + static_cast(params.numTokens * params.tokenDim * WorldSize * sizeof(__nv_bfloat16)), 0, 0, 0}; + rotateIfAllArrived(flag, rotationPending, arrivals, bytesToClear); + + // ==================== Residual + RMSNorm ==================== + __pipeline_wait_prior(0); + PackedBf16x8 residualIn; + residualIn.packed = storage.residual[threadIdx.x]; + hopEpilogue(params, flag, packedAccum, residualIn, threadOffset, rotationPending, arrivals, bytesToClear); +#endif +} + +// (CTA size, cluster size) of the one-shot grid for rows of tokenDim BF16 elements, eight per thread: the grid of the +// one-shot MNNVL all-reduce (its grid configuration with clusters and one access per scale group). Starts from the +// widest cluster that splits the row evenly, trades cluster width for CTAs of at least 128 threads, narrows the +// clusters while numTokens of them exceed the SMs, and widens them again down to 64-thread CTAs while the SMs allow. +std::pair hopGridConfig(int numTokens, int tokenDim, int smCount) +{ + int const threadsNeeded = (tokenDim + kEltsPerThread - 1) / kEltsPerThread; + int clusterSize = kMaxClusterSize; + while (threadsNeeded % clusterSize != 0 && clusterSize > 1) + { + clusterSize /= 2; + } + int const maxDivisibleClusterSize = clusterSize; + int blockSize = (threadsNeeded + clusterSize - 1) / clusterSize; + while (blockSize < 128 && clusterSize >= 2) + { + blockSize *= 2; + clusterSize /= 2; + } + while (numTokens * clusterSize > smCount && clusterSize > 1 && blockSize <= 512) + { + blockSize *= 2; + clusterSize /= 2; + } + while (clusterSize < maxDivisibleClusterSize) + { + int const candidateClusterSize = clusterSize * 2; + int const candidateBlockSize = (threadsNeeded + candidateClusterSize - 1) / candidateClusterSize; + if (candidateBlockSize < 64 || numTokens * candidateClusterSize > smCount) + { + break; + } + clusterSize = candidateClusterSize; + blockSize = candidateBlockSize; + } + return {blockSize, clusterSize}; +} + +// The launch attributes of both hops: PDL when enabled, the cluster, and the largest shared-memory carveout, so that a +// dependent launched early whose CTAs need most of the shared memory (the fused-A GEMM) can place them on the SMs that +// run this kernel's CTAs instead of starting them only when these exit. +struct HopLaunch +{ + cudaLaunchAttribute attrs[3]; + cudaLaunchConfig_t config; +}; + +void initHopLaunch(HopLaunch& launch, MnnvlHopParams const& params) +{ + TLLM_CHECK_WITH_INFO(common::getSMVersion() >= 90, "[GLM MNNVL hop] requires SM 90 or newer."); + TLLM_CHECK_WITH_INFO(mnnvlHopSupported(params.numTokens, params.tokenDim, params.nRanks), + "[GLM MNNVL hop] unsupported shape: %d tokens, hidden %d, %d ranks.", params.numTokens, params.tokenDim, + params.nRanks); + auto const [blockSize, clusterSize] + = hopGridConfig(params.numTokens, params.tokenDim, common::getMultiProcessorCount()); + launch.attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + launch.attrs[0].val.programmaticStreamSerializationAllowed = common::getEnvEnablePDL() ? 1 : 0; + launch.attrs[1].id = cudaLaunchAttributeClusterDimension; + launch.attrs[1].val.clusterDim.x = 1; + launch.attrs[1].val.clusterDim.y = clusterSize; + launch.attrs[1].val.clusterDim.z = 1; + launch.attrs[2].id = cudaLaunchAttributePreferredSharedMemoryCarveout; + launch.attrs[2].val.sharedMemCarveout = cudaSharedmemCarveoutMaxShared; + launch.config = {}; + launch.config.gridDim = dim3(params.numTokens, clusterSize, 1); + launch.config.blockDim = dim3(blockSize, 1, 1); + launch.config.dynamicSmemBytes = 0; + launch.config.stream = params.stream; + launch.config.attrs = launch.attrs; + launch.config.numAttrs = 3; +} + +HopKernelParams makeHopKernelParams(MnnvlHopParams const& params) +{ + HopKernelParams kernelParams{}; + kernelParams.normOut = static_cast<__nv_bfloat16*>(params.normOut); + kernelParams.residualOut = static_cast<__nv_bfloat16*>(params.residualOut); + kernelParams.shard = static_cast<__nv_bfloat16 const*>(params.input); + kernelParams.residualIn = static_cast<__nv_bfloat16 const*>(params.residualIn); + kernelParams.gamma = static_cast<__nv_bfloat16 const*>(params.gamma); + kernelParams.bufferPtrs = reinterpret_cast<__nv_bfloat16**>(params.bufferPtrsDev); + kernelParams.mcastPtr = static_cast<__nv_bfloat16*>(params.multicastPtr); + kernelParams.bufferFlags = params.bufferFlags; + kernelParams.numTokens = params.numTokens; + kernelParams.tokenDim = params.tokenDim; + kernelParams.rank = params.rank; + kernelParams.epsilon = params.epsilon; + return kernelParams; +} + +template +void launchHop(cudaLaunchConfig_t const& config, HopKernelParams const& kernelParams, bool shardPublished) +{ + if (shardPublished) + { + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, &mnnvlHopKernel, kernelParams)); + } + else + { + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, &mnnvlHopKernel, kernelParams)); + } +} + +template +void launchMoeFinalizeHop( + cudaLaunchConfig_t const& config, HopKernelParams const& kernelParams, MoeFinalizeKernelParams const& moeParams) +{ + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, &mnnvlMoeFinalizeHopKernel, kernelParams, moeParams)); +} + +} // namespace + +int64_t mnnvlHopBufferBytes(int numTokens, int tokenDim, int nRanks) +{ + return static_cast(numTokens) * tokenDim * nRanks * static_cast(sizeof(__nv_bfloat16)); +} + +bool mnnvlHopSupported(int numTokens, int tokenDim, int nRanks) +{ + if ((nRanks != 2 && nRanks != 4 && nRanks != 8 && nRanks != 16) || numTokens < 1 || tokenDim <= 0 + || tokenDim % kEltsPerThread != 0) + { + return false; + } + int64_t const bytes = mnnvlHopBufferBytes(numTokens, tokenDim, nRanks); + if (bytes > kOneShotBytes && !(nRanks == 16 && bytes <= kOneShotBytes16Ranks)) + { + return false; + } + auto const [blockSize, clusterSize] = hopGridConfig(numTokens, tokenDim, common::getMultiProcessorCount()); + return blockSize % kWarpSize == 0 && blockSize <= kMaxThreads && clusterSize > 1 && clusterSize <= kMaxClusterSize + && blockSize * clusterSize * kEltsPerThread == tokenDim; +} + +void mnnvlHop(MnnvlHopParams const& params) +{ + HopLaunch launch; + initHopLaunch(launch, params); + HopKernelParams const kernelParams = makeHopKernelParams(params); + switch (params.nRanks) + { + case 2: launchHop<2>(launch.config, kernelParams, params.shardPublished); break; + case 4: launchHop<4>(launch.config, kernelParams, params.shardPublished); break; + case 8: launchHop<8>(launch.config, kernelParams, params.shardPublished); break; + case 16: launchHop<16>(launch.config, kernelParams, params.shardPublished); break; + default: TLLM_CHECK_WITH_INFO(false, "[GLM MNNVL hop] unsupported world size %d.", params.nRanks); + } +} + +void mnnvlMoeFinalizeHop(MnnvlHopParams const& params, MnnvlMoeFinalizeParams const& moe) +{ + HopLaunch launch; + initHopLaunch(launch, params); + HopKernelParams const kernelParams = makeHopKernelParams(params); + MoeFinalizeKernelParams moeParams{}; + moeParams.fc2Output = static_cast<__nv_bfloat16 const*>(moe.fc2Output); + moeParams.expertWeights = static_cast<__nv_bfloat16 const*>(moe.expertWeights); + moeParams.expandedIdxToPermutedIdx = moe.expandedIdxToPermutedIdx; + moeParams.addend = static_cast<__nv_bfloat16 const*>(moe.addend); + moeParams.hiddenDimPadded = moe.hiddenDimPadded; + switch (params.nRanks) + { + case 2: launchMoeFinalizeHop<2>(launch.config, kernelParams, moeParams); break; + case 4: launchMoeFinalizeHop<4>(launch.config, kernelParams, moeParams); break; + case 8: launchMoeFinalizeHop<8>(launch.config, kernelParams, moeParams); break; + case 16: launchMoeFinalizeHop<16>(launch.config, kernelParams, moeParams); break; + default: TLLM_CHECK_WITH_INFO(false, "[GLM MNNVL MoE-finalize hop] unsupported world size %d.", params.nRanks); + } +} + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.h new file mode 100644 index 000000000000..693646e3eade --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.h @@ -0,0 +1,99 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +//! Routed experts per token of the MoE-finalize hop. +constexpr int kMnnvlHopMoeTopK = 8; + +//! \brief One BF16 one-shot MNNVL hop: the all-reduce of [numTokens, tokenDim] rows over nRanks ranks, plus the +//! residual and RMSNorm, on an MNNVL all-reduce workspace (three Lamport buffers and their flags, -0.0 as the empty +//! value; the workspace of the one-shot MNNVL all-reduce). +struct MnnvlHopParams +{ + int nRanks; //!< Ranks of the group: 2, 4, 8 or 16 + int rank; //!< This rank + int numTokens; //!< Rows + int tokenDim; //!< Elements per row + void** bufferPtrsDev; //!< Unicast addresses of every rank's workspace (a device array) + void* multicastPtr; //!< Multicast address of the workspace + uint32_t* bufferFlags; //!< Lamport flags of the workspace + void const* input; //!< This rank's partial rows [numTokens, tokenDim]; not read when shardPublished + void const* residualIn; //!< [numTokens, tokenDim] + void const* gamma; //!< RMSNorm weight [tokenDim] + float epsilon; //!< RMSNorm epsilon + void* residualOut; //!< Rows summed over the ranks plus residualIn [numTokens, tokenDim] + void* normOut; //!< RMSNorm of residualOut times gamma [numTokens, tokenDim] + bool shardPublished; //!< The producer of input already stored this rank's rows into the workspace + cudaStream_t stream; +}; + +//! \brief The routed experts' rows of the MoE-finalize hop: the do_finalize=false outputs of the MoE and the +//! shared-expert output added to them. +struct MnnvlMoeFinalizeParams +{ + void const* fc2Output; //!< FC2 rows [*, hiddenDimPadded], BF16 + int hiddenDimPadded; //!< Row pitch of fc2Output in elements + void const* expertWeights; //!< Routing weights [numTokens, kMnnvlHopMoeTopK], BF16 + int32_t const* expandedIdxToPermutedIdx; //!< FC2 row of each (token, k) [numTokens, kMnnvlHopMoeTopK]; -1 skips + void const* addend; //!< Added to the finalized rows [numTokens, tokenDim], BF16 +}; + +//! Whether the hops serve a BF16 [numTokens, tokenDim] all-reduce over nRanks ranks: 2, 4, 8 or 16 ranks, one-shot +//! (up to 1 MiB of Lamport data per rank, and up to 1.5 MiB on 16 ranks), and a grid of one row per cluster of 2 to 8 +//! CTAs of at most 128 threads (whole warps), eight elements per thread, that covers the row exactly. +bool mnnvlHopSupported(int numTokens, int tokenDim, int nRanks); + +//! Bytes of one Lamport buffer that a hop of numTokens x tokenDim BF16 elements over nRanks ranks uses. +int64_t mnnvlHopBufferBytes(int numTokens, int tokenDim, int nRanks); + +//! \brief The one-shot all-reduce with the residual add and RMSNorm (FastEpilogue, FastProtocol, one poll loop). +//! +//! Every rank stores its rows (-0.0 written as +0.0) through the multicast address at [token][rank][column] of the +//! current Lamport buffer, empties the dirty buffer, and polls all slots in one loop. The sum is in rank order in FP32, +//! from zero per chunk of 8 ranks, the chunks added in order, and rounded to BF16; the residual is added in BF16; the +//! sum of squares of the BF16 values goes through a transaction barrier of the cluster. The PDL dependents launch +//! right after the grid dependency wait, so they must read the outputs only after their own wait. +//! +//! With shardPublished the producer of input (the o_proj GEMM's Lamport publish) has stored this rank's rows into +//! every rank's current buffer: the kernel polls every slot, its own included, does not wait for its primary before +//! the exchange and reads the flags and the residual with volatile loads; it counts its CTA arrival after its poll +//! and waits for its primary at its end. The producer must trigger its dependents only after its own grid dependency +//! wait. +void mnnvlHop(MnnvlHopParams const& params); + +//! \brief The MoE finalize fused into the hop: each thread finalizes its eight elements of every routed row (the +//! token's FC2 rows weighted in k order with an FMA, -1 slots skipped, rounded to BF16, the addend added in FP32 and +//! the sum rounded once more) and publishes them as its rows; the rest is mnnvlHop without a published shard. +//! params.input and params.shardPublished are not read. Eight threads per CTA read the token's FC2 row indices and +//! routing weights before the grid dependency wait: the kernel that wrote them must have completed when this grid +//! launches (the MoE's FC2 triggers its dependents only after its own wait). +void mnnvlMoeFinalizeHop(MnnvlHopParams const& params, MnnvlMoeFinalizeParams const& moe); + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index a2d6c763597b..9113f3678eaa 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -135,6 +135,7 @@ add_th_common_family( allreduceOp.cpp alltoallOp.cpp asyncUlyssesOp.cpp + glmMnnvlHopOp.cpp ulyssesPostUnscatterOp.cpp ulyssesPermuteScatterOp.cpp helixPostProcessOp.cpp diff --git a/cpp/tensorrt_llm/thop/glmMnnvlHopOp.cpp b/cpp/tensorrt_llm/thop/glmMnnvlHopOp.cpp new file mode 100644 index 000000000000..dff31244e77e --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmMnnvlHopOp.cpp @@ -0,0 +1,217 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/mcastDevMemUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.h" +#include "tensorrt_llm/runtime/mcastDeviceMemory.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ +namespace +{ + +namespace glm = tensorrt_llm::kernels::glm_moe_dsa; + +bool isAlignedBf16(at::Tensor const& t) +{ + return t.scalar_type() == at::ScalarType::BFloat16 && t.is_contiguous() + && reinterpret_cast(t.data_ptr()) % 16 == 0; +} + +bool hopShapeSupported(int64_t numTokens, int64_t hiddenSize, int64_t nRanks) +{ + constexpr int64_t kIntMax = std::numeric_limits::max(); + if (numTokens > kIntMax || hiddenSize > kIntMax || nRanks > kIntMax) + { + return false; + } + return glm::mnnvlHopSupported(static_cast(numTokens), static_cast(hiddenSize), static_cast(nRanks)); +} + +void checkRowsAndGamma(at::Tensor const& residual, at::Tensor const& gamma, char const* op) +{ + TORCH_CHECK(residual.is_cuda() && residual.dim() == 2 && isAlignedBf16(residual), op, + ": residual must be a contiguous, 16-byte aligned BF16 [num_tokens, hidden] CUDA tensor"); + TORCH_CHECK(gamma.dim() == 1 && gamma.size(0) == residual.size(1) && isAlignedBf16(gamma), op, + ": gamma must be a contiguous, 16-byte aligned BF16 [hidden] tensor"); + TORCH_CHECK(gamma.device() == residual.device(), op, ": gamma must be on the device of the rows"); +} + +// The MNNVL all-reduce workspace that comm_buffer views ([3, N]: this rank's three Lamport buffers) for a hop of the +// rows of residual: the shape must be one the hops serve and its Lamport data must fit one buffer. +tensorrt_llm::runtime::McastDeviceMemory* hopWorkspace( + at::Tensor const& comm_buffer, at::Tensor const& buffer_flags, at::Tensor const& residual, char const* op) +{ + auto* mcastMem = tensorrt_llm::common::findMcastDevMemBuffer(comm_buffer.data_ptr()); + TORCH_CHECK(mcastMem != nullptr, op, ": comm_buffer must be obtained from a McastGPUBuffer instance"); + TORCH_CHECK(mcastMem->isMapped(), op, ": the MNNVL workspace handles are not attached"); + TORCH_CHECK(comm_buffer.dim() == 2 && comm_buffer.size(0) == 3 && comm_buffer.is_contiguous(), op, + ": comm_buffer must be the [3, N] view of the workspace's three Lamport buffers"); + TORCH_CHECK(buffer_flags.element_size() == 4 && buffer_flags.numel() >= 9 && buffer_flags.is_contiguous(), op, + ": buffer_flags must hold the workspace's nine 32-bit Lamport flags"); + TORCH_CHECK(comm_buffer.device() == residual.device() && buffer_flags.device() == residual.device(), op, + ": the workspace must be on the device of the rows"); + int64_t const numTokens = residual.size(0); + int64_t const hiddenSize = residual.size(1); + int64_t const nRanks = static_cast(mcastMem->getWorldSize()); + TORCH_CHECK(hopShapeSupported(numTokens, hiddenSize, nRanks), op, ": unsupported shape: ", numTokens, + " tokens, hidden ", hiddenSize, ", ", nRanks, " ranks"); + int64_t const bufferBytes = comm_buffer.size(1) * comm_buffer.element_size(); + int64_t const hopBytes = glm::mnnvlHopBufferBytes( + static_cast(numTokens), static_cast(hiddenSize), static_cast(nRanks)); + TORCH_CHECK(hopBytes <= bufferBytes, op, ": ", hopBytes, " bytes of Lamport data exceed the Lamport buffer of ", + bufferBytes, " bytes"); + return mcastMem; +} + +glm::MnnvlHopParams makeHopParams(tensorrt_llm::runtime::McastDeviceMemory* mcastMem, at::Tensor& buffer_flags, + at::Tensor const& residual, at::Tensor const& gamma, double epsilon, at::Tensor& normOut, at::Tensor& residualOut) +{ + glm::MnnvlHopParams params{}; + params.nRanks = static_cast(mcastMem->getWorldSize()); + params.rank = static_cast(mcastMem->getRank()); + params.numTokens = static_cast(residual.size(0)); + params.tokenDim = static_cast(residual.size(1)); + params.bufferPtrsDev = mcastMem->getBufferPtrsDev(); + params.multicastPtr = mcastMem->getMulticastPtr(); + params.bufferFlags = reinterpret_cast(buffer_flags.mutable_data_ptr()); + params.input = nullptr; + params.residualIn = residual.const_data_ptr(); + params.gamma = gamma.const_data_ptr(); + params.epsilon = static_cast(epsilon); + params.residualOut = residualOut.mutable_data_ptr(); + params.normOut = normOut.mutable_data_ptr(); + params.shardPublished = false; + params.stream = at::cuda::getCurrentCUDAStream(residual.get_device()); + return params; +} + +} // namespace + +// The one-shot MNNVL all-reduce of BF16 rows with the residual add and RMSNorm (kernels/glmMoeDsa/glmMnnvlHop.h) on +// the MNNVL all-reduce workspace comm_buffer / buffer_flags. Returns [norm_out, residual_out]. With shard_published the +// producer of input already stored this rank's rows into the workspace (the o_proj GEMM's Lamport publish), and input +// only gives the shape. +std::vector glmMnnvlHop(at::Tensor const& input, at::Tensor const& residual, at::Tensor const& gamma, + double epsilon, at::Tensor& comm_buffer, at::Tensor& buffer_flags, bool shard_published) +{ + char const* op = "glm_mnnvl_hop"; + checkRowsAndGamma(residual, gamma, op); + TORCH_CHECK(input.sizes() == residual.sizes(), op, ": input must have the shape of residual"); + if (!shard_published) + { + TORCH_CHECK(isAlignedBf16(input) && input.device() == residual.device(), op, + ": input must be a contiguous, 16-byte aligned BF16 tensor on the device of residual"); + } + at::Tensor normOut = at::empty_like(residual); + at::Tensor residualOut = at::empty_like(residual); + if (residual.size(0) == 0) + { + return {normOut, residualOut}; + } + auto* mcastMem = hopWorkspace(comm_buffer, buffer_flags, residual, op); + auto params = makeHopParams(mcastMem, buffer_flags, residual, gamma, epsilon, normOut, residualOut); + params.input = shard_published ? nullptr : input.const_data_ptr(); + params.shardPublished = shard_published; + glm::mnnvlHop(params); + return {normOut, residualOut}; +} + +// The TRTLLM-Gen MoE finalize (top-8, BF16) fused into the hop: the inputs are the MoE's do_finalize=false outputs (FC2 +// rows, routing weights, expanded-to-permuted index) and the addend (the shared experts' output). Returns [norm_out, +// residual_out], those of glm_mnnvl_hop on the finalized rows. The routing weights and index are read before the +// kernel's grid dependency wait: the kernel that wrote them must have completed when this one launches. +std::vector glmMnnvlMoeFinalizeHop(at::Tensor const& gemm2_output, at::Tensor const& expert_weights, + at::Tensor const& expanded_idx_to_permuted_idx, at::Tensor const& addend, at::Tensor const& residual, + at::Tensor const& gamma, double epsilon, at::Tensor& comm_buffer, at::Tensor& buffer_flags) +{ + char const* op = "glm_mnnvl_moe_finalize_hop"; + checkRowsAndGamma(residual, gamma, op); + int64_t const numTokens = residual.size(0); + TORCH_CHECK(gemm2_output.dim() == 2 && isAlignedBf16(gemm2_output), op, + ": gemm2_output must be a contiguous, 16-byte aligned BF16 2D tensor"); + TORCH_CHECK(gemm2_output.size(1) >= residual.size(1) && gemm2_output.size(1) % 8 == 0, op, + ": gemm2_output rows must hold the hidden dimension in a multiple of 8 elements"); + TORCH_CHECK(expanded_idx_to_permuted_idx.scalar_type() == at::ScalarType::Int + && expanded_idx_to_permuted_idx.dim() == 2 && expanded_idx_to_permuted_idx.is_contiguous() + && expanded_idx_to_permuted_idx.size(0) == numTokens + && expanded_idx_to_permuted_idx.size(1) == glm::kMnnvlHopMoeTopK, + op, ": expanded_idx_to_permuted_idx must be contiguous int32 [num_tokens, 8]"); + TORCH_CHECK(expert_weights.scalar_type() == at::ScalarType::BFloat16 && expert_weights.is_contiguous() + && expert_weights.sizes() == expanded_idx_to_permuted_idx.sizes(), + op, ": expert_weights must be contiguous BF16 of the index's shape"); + TORCH_CHECK(addend.sizes() == residual.sizes() && isAlignedBf16(addend), op, + ": addend must be a contiguous, 16-byte aligned BF16 tensor of the shape of residual"); + for (auto const* t : {&gemm2_output, &expert_weights, &expanded_idx_to_permuted_idx, &addend}) + { + TORCH_CHECK(t->device() == residual.device(), op, ": all inputs must be on one device"); + } + at::Tensor normOut = at::empty_like(residual); + at::Tensor residualOut = at::empty_like(residual); + if (numTokens == 0) + { + return {normOut, residualOut}; + } + auto* mcastMem = hopWorkspace(comm_buffer, buffer_flags, residual, op); + auto const params = makeHopParams(mcastMem, buffer_flags, residual, gamma, epsilon, normOut, residualOut); + glm::MnnvlMoeFinalizeParams moe{}; + moe.fc2Output = gemm2_output.const_data_ptr(); + moe.hiddenDimPadded = static_cast(gemm2_output.size(1)); + moe.expertWeights = expert_weights.const_data_ptr(); + moe.expandedIdxToPermutedIdx = expanded_idx_to_permuted_idx.const_data_ptr(); + moe.addend = addend.const_data_ptr(); + glm::mnnvlMoeFinalizeHop(params, moe); + return {normOut, residualOut}; +} + +// Whether glm_mnnvl_hop and glm_mnnvl_moe_finalize_hop serve [num_tokens, hidden_size] BF16 rows over num_ranks ranks +// on the current device. +bool glmMnnvlHopSupported(int64_t num_tokens, int64_t hidden_size, int64_t num_ranks) +{ + return hopShapeSupported(num_tokens, hidden_size, num_ranks); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_mnnvl_hop(Tensor input, Tensor residual, Tensor gamma, float epsilon, Tensor(a!) comm_buffer, " + "Tensor(b!) buffer_flags, bool shard_published) -> Tensor[]"); + m.def( + "glm_mnnvl_moe_finalize_hop(Tensor gemm2_output, Tensor expert_weights, Tensor expanded_idx_to_permuted_idx, " + "Tensor addend, Tensor residual, Tensor gamma, float epsilon, Tensor(a!) comm_buffer, Tensor(b!) buffer_flags) " + "-> Tensor[]"); + m.def("glm_mnnvl_hop_supported", &tensorrt_llm::torch_ext::glmMnnvlHopSupported); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_mnnvl_hop", &tensorrt_llm::torch_ext::glmMnnvlHop); + m.impl("glm_mnnvl_moe_finalize_hop", &tensorrt_llm::torch_ext::glmMnnvlMoeFinalizeHop); +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/_mnnvl_workspace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/_mnnvl_workspace.py new file mode 100644 index 000000000000..27d5f04504b3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/_mnnvl_workspace.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MNNVL Lamport workspaces of the target's hops. + +The hops (``mnnvl_hop``, ``mnnvl_moe_finalize_hop``) run on the MNNVL all-reduce workspace of the +engine's ``AllReduce`` (``get_or_scale_allreduce_mnnvl_workspace``): three Lamport buffers whose +empty value is -0.0, and nine uint32 flags (current buffer, dirty buffer, bytes per buffer, dirty +stages, bytes to clear per stage, arrival counter). The hops read, clear and rotate them as the +built-in MNNVL all-reduce does, so both can follow each other on one workspace. + +The workspace is replaced by a larger one when an all-reduce needs more than it holds, and a CUDA +graph keeps the workspace it launches on alive only when it was fetched during that graph's +capture. So the target fetches it for every hop it issues, eagerly and under capture, and never +keeps it across forwards. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch.distributed.ops import get_or_scale_allreduce_mnnvl_workspace +from tensorrt_llm.mapping import Mapping + +# Lamport buffers of an MNNVL all-reduce workspace. +NUM_LAMPORT_BUFFERS = 3 + + +def hop_buffer_bytes(num_tokens: int, hidden_size: int, tp_size: int) -> int: + """Bytes of one Lamport buffer that a hop of ``[num_tokens, hidden_size]`` BF16 rows uses.""" + return num_tokens * hidden_size * tp_size * torch.bfloat16.itemsize + + +def hop_supported(num_tokens: int, hidden_size: int, tp_size: int) -> bool: + """Whether the hops serve ``[num_tokens, hidden_size]`` BF16 rows over ``tp_size`` ranks. + + Asked of the current device: 2, 4, 8 or 16 ranks; one-shot (up to 1 MiB of Lamport data per + rank, 1.5 MiB on 16 ranks); one row per cluster of 2 to 8 CTAs of at most 128 threads, eight + elements per thread, covering the row exactly. Hidden 6144 on 16 ranks: 1 to 8 rows. Another + shape runs on the built-in ``AllReduce``. + """ + return bool(torch.ops.trtllm.glm_mnnvl_hop_supported(num_tokens, hidden_size, tp_size)) + + +def _all_reduce_workspace(mapping: Mapping, num_tokens: int, hidden_size: int) -> dict: + return get_or_scale_allreduce_mnnvl_workspace( + mapping, + torch.bfloat16, + buffer_size_bytes=hop_buffer_bytes(num_tokens, hidden_size, mapping.tp_size), + ) + + +def hop_workspace( + mapping: Mapping, num_tokens: int, hidden_size: int +) -> Tuple[torch.Tensor, torch.Tensor]: + """The workspace operands of a hop of ``[num_tokens, hidden_size]`` BF16 rows. + + Returns ``(comm_buffer, buffer_flags)``: the ``[3, N]`` BF16 view of this rank's three Lamport + buffers of the MNNVL all-reduce workspace of ``mapping``, and the workspace's nine uint32 + flags. The workspace grows first when its buffers hold less than the hop needs; it must not + grow while a CUDA graph is captured, so every shape runs once eagerly before its capture. + """ + workspace = _all_reduce_workspace(mapping, num_tokens, hidden_size) + comm_buffer = workspace["uc_buffer"].view(torch.bfloat16).view(NUM_LAMPORT_BUFFERS, -1) + return comm_buffer, workspace["buffer_flags"] + + +def publish_target( + mapping: Mapping, num_tokens: int, hidden_size: int +) -> Tuple[torch.Tensor, torch.Tensor]: + """Where the GEMM that produces a hop's input stores this rank's rows itself. + + Returns ``(lamport, flags)``: the multicast BF16 view (1-D) of the three Lamport buffers of the + workspace ``hop_workspace`` returns for the same shape, and its flags viewed as int32. The GEMM + (the o_proj Lamport publish) reads the current buffer index and the bytes per buffer from the + flags after its grid dependency wait, releases its dependents right after that wait, and stores + its BF16 output, -0.0 written as +0.0, in whole 16-byte vectors at ``[token][tp_rank][column]`` + of the current buffer. The hop that follows it on the stream then runs with + ``shard_published=True`` on that workspace. + """ + workspace = _all_reduce_workspace(mapping, num_tokens, hidden_size) + elements = NUM_LAMPORT_BUFFERS * workspace["buffer_size_bytes"] // torch.bfloat16.itemsize + lamport = workspace["handle"].get_mc_buffer((elements,), torch.bfloat16, 0) + return lamport, workspace["buffer_flags"].view(torch.int32) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_hop.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_hop.py new file mode 100644 index 000000000000..a9a591882939 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_hop.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One-shot MNNVL all-reduce with the residual add and RMSNorm, via the trtllm glm_mnnvl_hop op.""" + +from __future__ import annotations + +from typing import Sequence, Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._mnnvl_workspace import NUM_LAMPORT_BUFFERS + + +class GlmMnnvlHop(OpWrapper): + """All-reduce of BF16 rows over the ranks plus ``residual``, and the RMSNorm of that sum. + + ``residual_out = sum of input over the ranks + residual``, ``norm_out = RMSNorm(residual_out)``, + in one kernel on the MNNVL all-reduce workspace (``_mnnvl_workspace.hop_workspace``): every rank + stores its rows through the multicast address into its slot of the current Lamport buffer of + every rank and polls all slots in one loop. The sum is in FP32 in rank order (from zero per + chunk of 8 ranks, the chunks added in order) and rounded to BF16; the residual is added in + BF16 and the RMSNorm normalizes that BF16 sum, ``(x * rsqrt(mean(x^2) + epsilon)) * gamma``. + Every rank gets the same bits. Returns ``(norm_out, residual_out)``, both fresh. + + The kernel releases its PDL dependents right after its grid dependency wait (at its start when + the shard is published): a dependent reads ``norm_out`` and ``residual_out`` only after its + own grid dependency wait. + + With ``shard_published`` the GEMM that produced ``input`` (o_proj, through + ``_mnnvl_workspace.publish_target``) has already stored this rank's rows into the workspace; + ``input`` gives the shape only and is not read. That GEMM must be the hop's predecessor on the + stream and release its dependents only after its own grid dependency wait. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The c1 decode step at TP16: 8 rows on a verify step and draft step 0, 1 row on draft steps + # 1-6; every rank holds the full hidden size 6144. The world size is the test's. + CELLS: tuple[Cell, ...] = ( + Cell( + why="post-attention hop of the 78 verify layers and draft step 0: 8 rows, the shard " + "published by o_proj", + spec=dict(rows=8, hidden=6144, shard_published=True, zero_residual=False), + ), + Cell( + why="post-MLP hop of the dense layers 0-2: 8 rows, the shard in input", + spec=dict(rows=8, hidden=6144, shard_published=False, zero_residual=False), + ), + Cell( + why="draft step 0's eh_proj hop: input_layernorm with a zero residual, the sum as " + "residual_out", + spec=dict(rows=8, hidden=6144, shard_published=False, zero_residual=True), + ), + Cell( + why="post-attention hop of draft steps 1-6: one row, the shard published by o_proj", + spec=dict(rows=1, hidden=6144, shard_published=True, zero_residual=False), + ), + ) + + note = """ + Every rank of the group calls the hop in the same order, with the same shape and the same + `shard_published`; a disagreement hangs the group. `_mnnvl_workspace.hop_supported` states the + shapes the op takes (hidden 6144 on 16 ranks: 1 to 8 rows); it refuses others, which the target + runs on the built-in `AllReduce`. A -0.0 among the summed elements arrives as +0.0, because + -0.0 marks an empty Lamport slot. + """ + + def __call__( + self, + input: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + epsilon: float, + comm_buffer: torch.Tensor, + buffer_flags: torch.Tensor, + shard_published: bool, + ) -> Tuple[torch.Tensor, torch.Tensor]: + norm_out, residual_out = torch.ops.trtllm.glm_mnnvl_hop( + input, residual, gamma, epsilon, comm_buffer, buffer_flags, shard_published + ) + return norm_out, residual_out + + def reference( + self, + per_rank_inputs: Sequence[torch.Tensor], + residual: torch.Tensor, + gamma: torch.Tensor, + epsilon: float, + comm_buffer: torch.Tensor | None = None, + buffer_flags: torch.Tensor | None = None, + shard_published: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(norm_out, residual_out)`` in FP32, each rounded once. + + A collective: ``per_rank_inputs`` holds every rank's ``input`` in rank order, in place of + this rank's. ``residual`` is the calling rank's. ``comm_buffer``, ``buffer_flags`` and + ``shard_published`` decide how the rows travel, not what is computed, and are ignored. + """ + total = torch.zeros_like(residual, dtype=torch.float32) + for rows in per_rank_inputs: + total = total + rows.float() + summed = total + residual.float() + rms = torch.rsqrt(summed.pow(2).mean(-1, keepdim=True) + epsilon) + norm = summed * rms * gamma.float() + return norm.to(residual.dtype), summed.to(residual.dtype) + + def is_valid( + self, + input: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + epsilon: float, + comm_buffer: torch.Tensor, + buffer_flags: torch.Tensor, + shard_published: bool, + ) -> None: + # The kernel reads the slots as BF16 and takes -0.0 as the empty value: a workspace with + # another empty value (the 0xFFFFFFFF words of an exchange workspace) hands it stale words + # as rows. + assert comm_buffer.dtype == torch.bfloat16, ( + "comm_buffer must be the BF16 view of the MNNVL all-reduce workspace (hop_workspace)" + ) + # Flags of another workspace make the kernel poll, clear and rotate buffers whose state + # they do not describe: stale slots are summed, or a buffer still in use is emptied. + bytes_per_buffer = int(buffer_flags.view(torch.int32)[2].item()) + workspace_bytes = comm_buffer.numel() * comm_buffer.element_size() + assert NUM_LAMPORT_BUFFERS * bytes_per_buffer == workspace_bytes, ( + "buffer_flags must be the flags of comm_buffer's workspace" + ) + + +mnnvl_hop = GlmMnnvlHop() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_moe_finalize_hop.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_moe_finalize_hop.py new file mode 100644 index 000000000000..00bba97f581e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_moe_finalize_hop.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MoE finalize fused into the post-MoE MNNVL hop, via trtllm glm_mnnvl_moe_finalize_hop.""" + +from __future__ import annotations + +from typing import Sequence, Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._mnnvl_workspace import NUM_LAMPORT_BUFFERS + + +class GlmMnnvlMoeFinalizeHop(OpWrapper): + """``mnnvl_hop`` of the finalized MoE rows: the routed sum plus ``addend``, all-reduced. + + On each rank, row ``t`` of the finalized rows is ``addend[t] + sum over k of + expert_weights[t, k] * gemm2_output[expanded_idx_to_permuted_idx[t, k]]``, over the slots whose + index is not -1: the routed sum in FP32 with an FMA in ``k`` order, rounded to BF16, then the + addend added in FP32 and rounded once more, as the TRTLLM-Gen finalize with the shared + experts' output as addend computes it. The rows are then all-reduced with the residual add and + RMSNorm exactly as ``mnnvl_hop`` does (without a published shard), on the same workspace. + Returns ``(norm_out, residual_out)``, both fresh, the same bits on every rank. + + The kernel reads ``expert_weights`` and ``expanded_idx_to_permuted_idx`` before its grid + dependency wait, so the kernel that wrote them must have completed when this one launches: the + MoE's routing kernel, when this hop follows FC2 on the stream and FC2 releases its dependents + only after its own grid dependency wait; or a constant written outside graph capture. The + other operands are read after the wait. The PDL dependents launch right after that wait and + read the outputs only after their own. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The c1 decode step at TP16, hidden 6144, every expert local (moe_ep 1). + CELLS: tuple[Cell, ...] = ( + Cell( + why="post-MoE hop of the 75 MoE layers of a verify step: 8 rows, top-8 FC2 rows, the " + "shared experts' output as addend", + spec=dict(rows=8, hidden=6144, routing="top8"), + ), + Cell( + why="post-MoE hop of an MTP draft step: one row, its routed row in slot 0 with " + "weight 1 and slots 1-7 empty (-1), the shared experts' output as addend", + spec=dict(rows=1, hidden=6144, routing="routed_row"), + ), + ) + + note = """ + The shapes are those of `mnnvl_hop` (`_mnnvl_workspace.hop_supported`), with exactly 8 routing + slots per token; the same group discipline applies. FC2 rows may be wider than the hidden size + (`gemm2_output.shape[1]` is the row pitch); the first `hidden` columns are read. + """ + + def __call__( + self, + gemm2_output: torch.Tensor, + expert_weights: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + addend: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + epsilon: float, + comm_buffer: torch.Tensor, + buffer_flags: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + norm_out, residual_out = torch.ops.trtllm.glm_mnnvl_moe_finalize_hop( + gemm2_output, + expert_weights, + expanded_idx_to_permuted_idx, + addend, + residual, + gamma, + epsilon, + comm_buffer, + buffer_flags, + ) + return norm_out, residual_out + + def reference( + self, + per_rank_inputs: Sequence[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]], + residual: torch.Tensor, + gamma: torch.Tensor, + epsilon: float, + comm_buffer: torch.Tensor | None = None, + buffer_flags: torch.Tensor | None = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(norm_out, residual_out)`` in FP32, each rounded once. + + A collective: ``per_rank_inputs`` holds every rank's ``(gemm2_output, expert_weights, + expanded_idx_to_permuted_idx, addend)`` in rank order, in place of this rank's. + ``residual`` is the calling rank's. ``comm_buffer`` and ``buffer_flags`` are ignored. + """ + hidden = residual.shape[-1] + total = torch.zeros_like(residual, dtype=torch.float32) + for gemm2_output, expert_weights, permuted_idx, addend in per_rank_inputs: + taken = permuted_idx >= 0 + rows = gemm2_output[permuted_idx.clamp(min=0).long(), :hidden].float() + weights = torch.where(taken, expert_weights.float(), 0.0) + total = total + (weights.unsqueeze(-1) * rows).sum(dim=1) + addend.float() + summed = total + residual.float() + rms = torch.rsqrt(summed.pow(2).mean(-1, keepdim=True) + epsilon) + norm = summed * rms * gamma.float() + return norm.to(residual.dtype), summed.to(residual.dtype) + + def is_valid( + self, + gemm2_output: torch.Tensor, + expert_weights: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + addend: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + epsilon: float, + comm_buffer: torch.Tensor, + buffer_flags: torch.Tensor, + ) -> None: + # An index outside [-1, rows) reads memory past the FC2 rows as an expert row. + index = expanded_idx_to_permuted_idx + in_range = (index >= -1) & (index < gemm2_output.shape[0]) + assert bool(in_range.all()), ( + "expanded_idx_to_permuted_idx must index gemm2_output's rows, or be -1" + ) + # As for mnnvl_hop: a workspace with another empty value, or flags of another workspace, + # make the kernel sum stale slots or empty a buffer still in use. + assert comm_buffer.dtype == torch.bfloat16, ( + "comm_buffer must be the BF16 view of the MNNVL all-reduce workspace (hop_workspace)" + ) + bytes_per_buffer = int(buffer_flags.view(torch.int32)[2].item()) + workspace_bytes = comm_buffer.numel() * comm_buffer.element_size() + assert NUM_LAMPORT_BUFFERS * bytes_per_buffer == workspace_bytes, ( + "buffer_flags must be the flags of comm_buffer's workspace" + ) + + +mnnvl_moe_finalize_hop = GlmMnnvlMoeFinalizeHop() diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_hop_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_hop_op_matrix.py new file mode 100644 index 000000000000..f8847b2712f7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_hop_op_matrix.py @@ -0,0 +1,420 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for the GLM-5.2 target's mnnvl_hop entry. + +A collective cannot be exercised in one process, so this script is its own launcher (see +``_mnnvl_ranks``): run plainly it spawns one rank per device in CUDA_VISIBLE_DEVICES, + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _mnnvl_hop_op_matrix.py + +and ``--rank-worker`` runs one rank of a world started elsewhere (the 16-rank c1 world). The +checks are one fixed sequence on one MNNVL all-reduce workspace, which each check leaves rotated +for the next; collected as tests they would run at world size 1, so neither this file's name nor +its functions match pytest's patterns. The collected entry point is +``test_modeling_v2_mnnvl_hop_op_matrix.py`` beside this file. + +Beyond the entry's cells: exact sums, the Lamport flags, -0.0 rows, the published shard, +interleaving with the built-in MNNVL all-reduce on its workspace, CUDA graph replays, ranks +arriving late, and the guards. On a world whose one-shot all-reduce stays under 1 MiB the built-in +op computes the same values in the same order, so the hop matches it bit for bit. +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from typing import Any, List, Tuple + +import _mnnvl_ranks as ranks +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "mnnvl_hop requires CUDA devices" + +DEADLINE_S = 1200 +HIDDEN = 6144 +EPS = 1e-5 + +# Bound in the rank body: the entry and the workspace helpers import tensorrt_llm. +entry: Any = None +ws: Any = None +validating: Any = None +fusion_op: Any = None +RANK = 0 +WORLD = 1 +MAPPING: Any = None + + +def _rows(rows: int, seed: int, rank: int) -> torch.Tensor: + """One rank's partial rows: their sum over the ranks has unit variance.""" + return ranks.normal((rows, HIDDEN), seed * 1009 + rank, scale=WORLD**-0.5) + + +def _residual(rows: int, seed: int, zero: bool = False) -> torch.Tensor: + """The residual stream: the same on every rank, as on the target.""" + if zero: + return torch.zeros((rows, HIDDEN), dtype=torch.bfloat16, device="cuda") + return ranks.normal((rows, HIDDEN), seed * 1009 + 500) + + +def _gamma(seed: int) -> torch.Tensor: + """An RMSNorm weight near 1, the same on every rank.""" + return (1.0 + ranks.normal(HIDDEN, seed, scale=0.1, dtype=torch.float32)).to(torch.bfloat16) + + +def _compare(norm, residual_out, ref_norm, ref_residual, where: str) -> None: + """Bands of the hop against the FP32 reference. + + residual_out: the kernel rounds the FP32 sum over the ranks to BF16 and adds the residual in + BF16, two roundings of half an ulp, each of a value at most the row's largest (inputs scaled + so that the sum and the residual are of one size): at most about 2 ulp of the row scale per + element, 3 allowed; about 0.6 ulp of relative RMS, 1.5 allowed. + + norm_out: residual_out's error carried through the scale (gamma within 0.6-1.4), the sum of + squares of the BF16 values (each square rounded to BF16, below 1/4 ulp on the scale) and the + output rounding (half an ulp): at most about 3.5 ulp per element, 5 allowed; about 0.7 ulp of + relative RMS, 1.5 allowed. Derived, not fitted: a misplaced slot or rank misses by orders. + """ + from _validating import assert_within_ulp + + try: + assert_within_ulp(residual_out, ref_residual, element_ulp=3.0, rms_ulp=1.5) + assert_within_ulp(norm, ref_norm, element_ulp=5.0, rms_ulp=1.5) + except AssertionError as exc: + raise AssertionError(f"{where}: {exc}") from exc + + +def _builtin_one_shot(rows: int) -> bool: + """Whether the built-in MNNVL all-reduce runs one-shot here, in the hops' reduction order.""" + return rows * HIDDEN * WORLD * torch.bfloat16.itemsize <= ranks.BUILTIN_ONE_SHOT_BYTES + + +def _builtin_hop(x, residual, gamma, comm_buffer, buffer_flags) -> Tuple[torch.Tensor, ...]: + """The built-in MNNVL residual + RMSNorm all-reduce on the same workspace.""" + norm, residual_out = torch.ops.trtllm.mnnvl_fusion_allreduce( + x, gamma, residual, EPS, comm_buffer, buffer_flags, True, None, fusion_op + ) + return norm, residual_out + + +def _word_offsets(rows: int) -> torch.Tensor: + """Offsets, in 32-bit words of a Lamport buffer, of this rank's slot of every row.""" + tokens = torch.arange(rows, device="cuda").view(-1, 1) + pairs = torch.arange(0, HIDDEN, 2, device="cuda").view(1, -1) + return (((tokens * WORLD + RANK) * HIDDEN + pairs) // 2).view(-1) + + +def _publish(x: torch.Tensor, offsets: torch.Tensor | None = None) -> None: + """Store this rank's rows into every rank's current Lamport buffer, as the o_proj publish does. + + The GEMM's epilogue reads the current buffer from the flags and stores its BF16 output, -0.0 + as +0.0, at [token][rank][column]. Done here with 32-bit stores through the multicast view: + a word is the unit the hop's poll tests, so a half-written slot is never taken for a full one. + The index comes from the flags on the device, so this also runs inside a CUDA graph. + """ + rows = x.shape[0] + lamport, flags = ws.publish_target(MAPPING, rows, HIDDEN) + words = lamport.view(torch.int32) + if offsets is None: + offsets = _word_offsets(rows) + base = flags[0].long() * (flags[2].long() // 4) + payload = torch.where(x == 0, torch.zeros_like(x), x).contiguous().view(torch.int32) + words.index_copy_(0, base + offsets, payload.view(-1)) + + +def _hop(x, residual, gamma, shard_published: bool) -> Tuple[torch.Tensor, torch.Tensor]: + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, x.shape[0], HIDDEN) + if shard_published: + _publish(x) + return entry.mnnvl_hop(x, residual, gamma, EPS, comm_buffer, buffer_flags, shard_published) + + +def check_certified_cells() -> None: + """Every cell: the band against the reference, the same bits on every rank, and the bits of + the built-in one-shot all-reduce where it runs one-shot.""" + gamma = _gamma(7) + for i, cell in enumerate(entry.mnnvl_hop.CELLS): + spec = cell.spec + rows, published, zero = spec["rows"], spec["shard_published"], spec["zero_residual"] + assert spec["hidden"] == HIDDEN, cell.why + seed = 100 + 13 * i + per_rank = [_rows(rows, seed, r) for r in range(WORLD)] + x, residual = per_rank[RANK], _residual(rows, seed, zero) + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, rows, HIDDEN) + if published: + _publish(x) + with validating(entry.mnnvl_hop): + norm, residual_out = entry.mnnvl_hop( + x, residual, gamma, EPS, comm_buffer, buffer_flags, published + ) + ref_norm, ref_residual = entry.mnnvl_hop.reference(per_rank, residual, gamma, EPS) + _compare(norm, residual_out, ref_norm, ref_residual, cell.why) + ranks.assert_same_on_every_rank(norm, f"{cell.why}: norm_out") + ranks.assert_same_on_every_rank(residual_out, f"{cell.why}: residual_out") + if _builtin_one_shot(rows): + builtin_norm, builtin_residual = _builtin_hop( + x, residual, gamma, comm_buffer, buffer_flags + ) + ranks.assert_bitwise(norm, builtin_norm, f"{cell.why}: norm_out vs built-in") + ranks.assert_bitwise(residual_out, builtin_residual, f"{cell.why}: residual_out") + dist.barrier() + + +def check_exact_sums_are_exact() -> None: + """Rows on the 1/8 grid sum exactly: residual_out equals the reference bit for bit. + + Any slot read from the wrong rank, row or column, or summed twice, moves an exact result. + """ + gamma = _gamma(11) + for rows, published in ((8, False), (8, True), (1, True), (3, False)): + seed = 2000 + rows + 7 * published + # |row| <= 1 per rank, |sum| <= 16 on 16 ranks, |residual| <= 8: all exact in BF16. + per_rank = [ranks.eighths((rows, HIDDEN), seed * 31 + r, bound=8) for r in range(WORLD)] + residual = ranks.eighths((rows, HIDDEN), seed * 31 + 97, bound=64) + norm, residual_out = _hop(per_rank[RANK], residual, gamma, published) + ref_norm, ref_residual = entry.mnnvl_hop.reference(per_rank, residual, gamma, EPS) + where = f"exact rows={rows} published={published}" + ranks.assert_bitwise(residual_out, ref_residual, where) + _compare(norm, residual_out, ref_norm, ref_residual, where) + dist.barrier() + + +def check_lamport_flags_rotate() -> None: + """After each hop the current buffer is the dirty one, with the hop's bytes to clear.""" + gamma = _gamma(13) + for step in range(4): + rows = (8, 1, 5, 8)[step] + published = step % 2 == 1 + _, buffer_flags = ws.hop_workspace(MAPPING, rows, HIDDEN) + before = ranks.lamport_flags(buffer_flags) + _hop(_rows(rows, 3000 + step, RANK), _residual(rows, 3000 + step), gamma, published) + after = ranks.lamport_flags(buffer_flags) + bytes_to_clear = rows * HIDDEN * WORLD * torch.bfloat16.itemsize + ranks.assert_rotated(before, after, bytes_to_clear, f"step {step}") + dist.barrier() + + +def check_negative_zero_rows_arrive_as_positive_zero() -> None: + """-0.0 marks an empty slot; rows of -0.0 must neither hang the exchange nor sum to -0.0.""" + gamma = _gamma(17) + for published in (False, True): + x = torch.full((8, HIDDEN), -0.0, dtype=torch.bfloat16, device="cuda") + residual = ranks.eighths((8, HIDDEN), 4000, bound=64) + # A residual of -0.0 plus a sum of +0.0 is +0.0, and any other residual is itself. + residual = torch.where(residual == 0, torch.ones_like(residual), residual) + _, residual_out = _hop(x, residual, gamma, published) + ranks.assert_bitwise(residual_out, residual, f"-0.0 rows published={published}") + dist.barrier() + + +def check_interleaved_with_the_builtin_all_reduce() -> None: + """The hops and the built-in MNNVL all-reduce (one-shot and two-shot) share one workspace. + + They read, clear and rotate the same flags; every result stays right whichever kernel ran + before it. + """ + gamma = _gamma(19) + for round_ in range(3): + seed = 5000 + 100 * round_ + for kind, rows in ( + ("hop", 8), + ("builtin", 2), + ("published", 1), + ("builtin_two_shot", 64), + ("published", 8), + ("hop", 1), + ): + per_rank = [_rows(rows, seed + rows, r) for r in range(WORLD)] + residual = _residual(rows, seed + rows) + if kind.startswith("builtin"): + # 64 rows are two-shot for the built-in all-reduce; a workspace that holds the + # rows one-shot holds them two-shot too. + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, rows, HIDDEN) + norm, residual_out = _builtin_hop( + per_rank[RANK], residual, gamma, comm_buffer, buffer_flags + ) + else: + norm, residual_out = _hop(per_rank[RANK], residual, gamma, kind == "published") + ref_norm, ref_residual = entry.mnnvl_hop.reference(per_rank, residual, gamma, EPS) + _compare(norm, residual_out, ref_norm, ref_residual, f"round {round_} {kind} {rows}") + dist.barrier() + + +def check_cuda_graph_replays() -> None: + """A captured decode-step chain replays right: publish, published hop, hop, built-in hop. + + The flags rotate on the device, so every replay runs on the buffers the previous kernel left + current; the published rows land where the flags say at replay time. + """ + gamma = _gamma(23) + rows = 8 + x_static = _rows(rows, 6000, RANK).clone() + residual_static = _residual(rows, 6000).clone() + x2_static = _rows(1, 6001, RANK).clone() + offsets = _word_offsets(rows) + + def chain(): + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, rows, HIDDEN) + _publish(x_static, offsets) + norm1, res1 = entry.mnnvl_hop( + x_static, residual_static, gamma, EPS, comm_buffer, buffer_flags, True + ) + norm2, res2 = entry.mnnvl_hop(norm1, res1, gamma, EPS, comm_buffer, buffer_flags, False) + comm_buffer1, buffer_flags1 = ws.hop_workspace(MAPPING, 1, HIDDEN) + norm3, res3 = entry.mnnvl_hop( + x2_static, res2[:1].contiguous(), gamma, EPS, comm_buffer1, buffer_flags1, False + ) + return norm1, res1, norm2, res2, norm3, res3 + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + chain() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs = chain() + for replay in range(7): + seed = 6100 + 10 * replay + per_rank = [_rows(rows, seed, r) for r in range(WORLD)] + per_rank1 = [_rows(1, seed + 1, r) for r in range(WORLD)] + x_static.copy_(per_rank[RANK]) + residual_static.copy_(_residual(rows, seed)) + x2_static.copy_(per_rank1[RANK]) + torch.cuda.synchronize() + dist.barrier() + graph.replay() + torch.cuda.synchronize() + norm1, res1, norm2, res2, norm3, res3 = outputs + ref = entry.mnnvl_hop.reference(per_rank, residual_static, gamma, EPS) + _compare(norm1, res1, *ref, f"replay {replay} published hop") + # The second hop sums norm1 over the ranks: build its inputs from every rank's norm1. + gathered = [torch.empty_like(norm1) for _ in range(WORLD)] + dist.all_gather(gathered, norm1) + ref = entry.mnnvl_hop.reference(gathered, res1, gamma, EPS) + _compare(norm2, res2, *ref, f"replay {replay} hop") + ref = entry.mnnvl_hop.reference(per_rank1, res2[:1], gamma, EPS) + _compare(norm3, res3, *ref, f"replay {replay} one-row hop") + ranks.assert_same_on_every_rank(norm3, f"replay {replay} one-row hop") + dist.barrier() + del graph + torch.cuda.synchronize() + + +def check_late_ranks() -> None: + """Ranks reach the hops at different times; back-to-back hops still exchange the right rows. + + Rank r sleeps on its stream before each hop, so the early ranks poll while the late ones + have not published, and a fast rank's next hop publishes into a buffer a slow rank has only + just emptied. + """ + gamma = _gamma(29) + issued: List[Tuple[Any, ...]] = [] + for step in range(6): + rows = (8, 1)[step % 2] + published = step % 3 == 0 + seed = 7000 + step + per_rank = [_rows(rows, seed, r) for r in range(WORLD)] + residual = _residual(rows, seed) + torch.cuda._sleep(((RANK * 7 + step * 3) % WORLD) * 400_000) + issued.append((per_rank, residual, *_hop(per_rank[RANK], residual, gamma, published))) + torch.cuda.synchronize() + for step, (per_rank, residual, norm, residual_out) in enumerate(issued): + ref_norm, ref_residual = entry.mnnvl_hop.reference(per_rank, residual, gamma, EPS) + _compare(norm, residual_out, ref_norm, ref_residual, f"late ranks step {step}") + + +def check_the_op_refuses_unsupported_shapes() -> None: + """Shapes outside hop_supported are refused loudly, before any rank enters the exchange.""" + assert ws.hop_supported(8, HIDDEN, WORLD) + assert ws.hop_supported(1, HIDDEN, WORLD) + # On 16 ranks: 8 rows of 6144 are the 1.5 MiB one-shot limit; 9 rows are above it. + assert ws.hop_supported(8, HIDDEN, 16) and not ws.hop_supported(9, HIDDEN, 16) + # 6152 columns are 769 eight-element vectors: no cluster of 2 or more CTAs splits them. + assert not ws.hop_supported(1, HIDDEN + 8, WORLD) + assert not ws.hop_supported(1, HIDDEN, 32) + too_many = 1 + while ws.hop_supported(too_many, HIDDEN, WORLD): + too_many += 1 + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, 1, HIDDEN) + for rows, hidden in ((too_many, HIDDEN), (1, HIDDEN + 8)): + x = torch.zeros((rows, hidden), dtype=torch.bfloat16, device="cuda") + g = torch.ones(hidden, dtype=torch.bfloat16, device="cuda") + try: + entry.mnnvl_hop(x, x, g, EPS, comm_buffer, buffer_flags, False) + except RuntimeError as exc: + assert "unsupported shape" in str(exc), exc + else: + raise AssertionError(f"the op took {rows} x {hidden} rows") + + +def check_the_guard_refuses_flags_of_another_workspace() -> None: + """Flags that do not describe comm_buffer's buffers are refused by is_valid.""" + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, 8, HIDDEN) + foreign = buffer_flags.clone() + foreign.view(torch.int32)[2] += 16 + x = _rows(8, 8000, RANK) + try: + with validating(entry.mnnvl_hop): + entry.mnnvl_hop(x, x, _gamma(37), EPS, comm_buffer, foreign, False) + except AssertionError as exc: + assert "flags of comm_buffer's workspace" in str(exc), exc + else: + raise AssertionError("is_valid admitted foreign flags") + + +def check_the_guard_refuses_a_workspace_of_another_empty_value() -> None: + """A workspace view that is not BF16 (an exchange workspace's words) is refused.""" + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, 8, HIDDEN) + x = _rows(8, 8100, RANK) + try: + with validating(entry.mnnvl_hop): + entry.mnnvl_hop( + x, x, _gamma(41), EPS, comm_buffer.view(torch.int32), buffer_flags, False + ) + except AssertionError as exc: + assert "BF16 view" in str(exc), exc + else: + raise AssertionError("is_valid admitted a non-BF16 workspace view") + + +CHECKS = ( + check_certified_cells, + check_exact_sums_are_exact, + check_lamport_flags_rotate, + check_negative_zero_rows_arrive_as_positive_zero, + check_interleaved_with_the_builtin_all_reduce, + check_cuda_graph_replays, + check_late_ranks, + check_the_op_refuses_unsupported_shapes, + check_the_guard_refuses_flags_of_another_workspace, + check_the_guard_refuses_a_workspace_of_another_empty_value, +) + + +def _run_one_rank() -> None: + global entry, ws, validating, fusion_op, RANK, WORLD, MAPPING + RANK, WORLD, MAPPING = ranks.init_rank() + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.comm import ( # noqa: E501 + _mnnvl_workspace, + mnnvl_hop, + ) + from tensorrt_llm.functional import AllReduceFusionOp + + entry, ws = mnnvl_hop, _mnnvl_workspace + fusion_op = int(AllReduceFusionOp.RESIDUAL_RMS_NORM) + # Run as a script, not collected: the shared _validating module is reached the plain way. + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating as _validating_cm + + validating = _validating_cm + ranks.run_checks(CHECKS, RANK) + + +if __name__ == "__main__": + if ranks.WORKER_FLAG in sys.argv: + _run_one_rank() + ranks.spawn_ranks(Path(__file__).resolve(), DEADLINE_S) + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_moe_finalize_hop_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_moe_finalize_hop_op_matrix.py new file mode 100644 index 000000000000..2544435a37d6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_moe_finalize_hop_op_matrix.py @@ -0,0 +1,354 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for the GLM-5.2 target's mnnvl_moe_finalize_hop entry. + +Its own launcher, like ``_mnnvl_hop_op_matrix.py`` (see ``_mnnvl_ranks``): + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _mnnvl_moe_finalize_hop_op_matrix.py + +The fused kernel must equal the two steps it replaces bit for bit: the TRTLLM-Gen finalize with +the shared experts' output as addend (computed here exactly in torch: a BF16 weight times a BF16 +row is exact in FP32, so the separate multiply and add are the FMA), followed by ``mnnvl_hop`` on +its rows. Beyond the cells: empty (-1) slots, FC2 rows wider than the hidden size, interleaving with +the plain hops on one workspace, CUDA graph replays, and the index guard. The collected entry +point is ``test_modeling_v2_mnnvl_moe_finalize_hop_op_matrix.py``. +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from typing import Any, List, Tuple + +import _mnnvl_ranks as ranks +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "mnnvl_moe_finalize_hop requires CUDA devices" + +DEADLINE_S = 1200 +HIDDEN = 6144 +EPS = 1e-5 +TOP_K = 8 + +# Bound in the rank body: the entries and the workspace helpers import tensorrt_llm. +entry: Any = None +hop: Any = None +ws: Any = None +validating: Any = None +RANK = 0 +WORLD = 1 +MAPPING: Any = None + +MoeInputs = Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] + + +def _moe_inputs( + rows: int, seed: int, rank: int, routing: str, pitch: int = HIDDEN, empty_slots: bool = False +) -> MoeInputs: + """One rank's ``(gemm2_output, expert_weights, expanded_idx_to_permuted_idx, addend)``. + + ``top8``: rows * 8 FC2 rows plus 5 padding rows, a random permutation of them over the + (token, k) slots, positive routing weights; with ``empty_slots`` some slots are -1. + ``routed_row``: the MTP layer's add in the hop, one FC2 row per token (its routed output) in + slot 0 with weight 1, slots 1-7 empty. Scaled so that the finalized rows summed over the ranks + have unit variance. + """ + scale = WORLD**-0.5 + base = seed * 7919 + rank * 131 + addend = ranks.normal((rows, HIDDEN), base + 1, scale=0.5 * scale) + if routing == "routed_row": + gemm2_output = ranks.normal((rows, pitch), base + 2, scale=0.5 * scale) + weights = torch.zeros((rows, TOP_K), dtype=torch.bfloat16, device="cuda") + weights[:, 0] = 1.0 + index = torch.full((rows, TOP_K), -1, dtype=torch.int32, device="cuda") + index[:, 0] = torch.arange(rows, dtype=torch.int32, device="cuda") + return gemm2_output, weights, index, addend + assert routing == "top8", routing + fc2_rows = rows * TOP_K + 5 + gemm2_output = ranks.normal((fc2_rows, pitch), base + 3, scale=scale) + gen = torch.Generator(device="cuda").manual_seed(base + 4) + index = torch.randperm(fc2_rows, generator=gen, device="cuda")[: rows * TOP_K] + index = index.to(torch.int32).view(rows, TOP_K).contiguous() + if empty_slots: + drop = torch.rand((rows, TOP_K), generator=gen, device="cuda") < 0.3 + index = torch.where(drop, torch.full_like(index, -1), index) + weights = (torch.rand((rows, TOP_K), generator=gen, device="cuda") * 0.25).to(torch.bfloat16) + return gemm2_output, weights, index, addend + + +def _finalized_rows(inputs: MoeInputs) -> torch.Tensor: + """The TRTLLM-Gen finalize with an addend, bit for bit: the routed sum in FP32 in k order + (each weight times row exact in FP32, so the add is the FMA's rounding), rounded to BF16, the + addend added in FP32 and rounded once more.""" + gemm2_output, weights, index, addend = inputs + routed = torch.zeros(addend.shape, dtype=torch.float32, device="cuda") + for k in range(TOP_K): + slot = index[:, k] + taken = (slot >= 0).view(-1, 1) + rows = gemm2_output[slot.clamp(min=0).long(), :HIDDEN].float() + product = weights[:, k].float().view(-1, 1) * rows + routed = torch.where(taken, routed + product, routed) + return (addend.float() + routed.to(torch.bfloat16).float()).to(torch.bfloat16) + + +def _residual(rows: int, seed: int) -> torch.Tensor: + """The residual stream: the same on every rank, as on the target.""" + return ranks.normal((rows, HIDDEN), seed * 1009 + 500) + + +def _gamma(seed: int) -> torch.Tensor: + return (1.0 + ranks.normal(HIDDEN, seed, scale=0.1, dtype=torch.float32)).to(torch.bfloat16) + + +def _compare(norm, residual_out, ref_norm, ref_residual, where: str) -> None: + """Bands of the fused hop against the FP32 reference. + + On top of mnnvl_hop's two roundings (the sum, the residual add) each rank rounds its routed + sum and then the addend's sum to BF16 before the exchange: two more half-ulp roundings of a + rank's row, whose sum over the ranks (independent signs) stays near one ulp of the summed + row's scale. residual_out: 4 ulp per element and 2 ulp of relative RMS allowed; norm_out: 6 + and 2. Derived, not fitted; the fused-vs-two-step gate in the checks is exact. + """ + from _validating import assert_within_ulp + + try: + assert_within_ulp(residual_out, ref_residual, element_ulp=4.0, rms_ulp=2.0) + assert_within_ulp(norm, ref_norm, element_ulp=6.0, rms_ulp=2.0) + except AssertionError as exc: + raise AssertionError(f"{where}: {exc}") from exc + + +def _fused(inputs: MoeInputs, residual, gamma) -> Tuple[torch.Tensor, torch.Tensor]: + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, residual.shape[0], HIDDEN) + gemm2_output, weights, index, addend = inputs + return entry.mnnvl_moe_finalize_hop( + gemm2_output, weights, index, addend, residual, gamma, EPS, comm_buffer, buffer_flags + ) + + +def _two_step(inputs: MoeInputs, residual, gamma) -> Tuple[torch.Tensor, torch.Tensor]: + """The finalize, then the plain hop on its rows.""" + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, residual.shape[0], HIDDEN) + rows = _finalized_rows(inputs) + return hop.mnnvl_hop(rows, residual, gamma, EPS, comm_buffer, buffer_flags, False) + + +def check_certified_cells() -> None: + """Every cell: equal to finalize + hop bit for bit, within the band of the reference, and the + same on every rank.""" + gamma = _gamma(7) + for i, cell in enumerate(entry.mnnvl_moe_finalize_hop.CELLS): + spec = cell.spec + rows = spec["rows"] + assert spec["hidden"] == HIDDEN, cell.why + seed = 100 + 13 * i + per_rank = [_moe_inputs(rows, seed, r, spec["routing"]) for r in range(WORLD)] + residual = _residual(rows, seed) + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, rows, HIDDEN) + gemm2_output, weights, index, addend = per_rank[RANK] + with validating(entry.mnnvl_moe_finalize_hop): + norm, residual_out = entry.mnnvl_moe_finalize_hop( + gemm2_output, + weights, + index, + addend, + residual, + gamma, + EPS, + comm_buffer, + buffer_flags, + ) + ref_norm, ref_residual = entry.mnnvl_moe_finalize_hop.reference( + per_rank, residual, gamma, EPS + ) + _compare(norm, residual_out, ref_norm, ref_residual, cell.why) + two_step_norm, two_step_residual = _two_step(per_rank[RANK], residual, gamma) + ranks.assert_bitwise(norm, two_step_norm, f"{cell.why}: norm_out vs finalize + hop") + ranks.assert_bitwise(residual_out, two_step_residual, f"{cell.why}: residual_out") + ranks.assert_same_on_every_rank(norm, f"{cell.why}: norm_out") + ranks.assert_same_on_every_rank(residual_out, f"{cell.why}: residual_out") + dist.barrier() + + +def check_empty_slots_and_wide_rows() -> None: + """-1 slots are skipped and FC2 rows wider than the hidden size are read at their pitch.""" + gamma = _gamma(11) + for rows, pitch, empty in ((8, HIDDEN, True), (8, HIDDEN + 128, False), (3, HIDDEN + 64, True)): + seed = 2000 + rows + pitch + per_rank = [ + _moe_inputs(rows, seed, r, "top8", pitch=pitch, empty_slots=empty) + for r in range(WORLD) + ] + residual = _residual(rows, seed) + norm, residual_out = _fused(per_rank[RANK], residual, gamma) + two_step = _two_step(per_rank[RANK], residual, gamma) + where = f"rows={rows} pitch={pitch} empty slots={empty}" + ranks.assert_bitwise(norm, two_step[0], f"{where}: norm_out vs finalize + hop") + ranks.assert_bitwise(residual_out, two_step[1], f"{where}: residual_out") + ref = entry.mnnvl_moe_finalize_hop.reference(per_rank, residual, gamma, EPS) + _compare(norm, residual_out, *ref, where) + dist.barrier() + + +def check_lamport_flags_rotate() -> None: + """The fused hop rotates the workspace's flags as the plain hop does.""" + gamma = _gamma(13) + for step, rows in enumerate((8, 1, 8)): + _, buffer_flags = ws.hop_workspace(MAPPING, rows, HIDDEN) + before = ranks.lamport_flags(buffer_flags) + routing = "top8" if rows > 1 else "routed_row" + _fused(_moe_inputs(rows, 3000 + step, RANK, routing), _residual(rows, step), gamma) + after = ranks.lamport_flags(buffer_flags) + bytes_to_clear = rows * HIDDEN * WORLD * torch.bfloat16.itemsize + ranks.assert_rotated(before, after, bytes_to_clear, f"step {step}") + dist.barrier() + + +def check_interleaved_with_the_hops() -> None: + """Fused hops, plain hops and the built-in all-reduce follow each other on one workspace.""" + from tensorrt_llm.functional import AllReduceFusionOp + + gamma = _gamma(17) + issued: List[Tuple[Any, ...]] = [] + for step in range(9): + rows = (8, 1, 8)[step % 3] + seed = 4000 + step + residual = _residual(rows, seed) + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, rows, HIDDEN) + if step % 3 == 2: + per_rank_rows = [ranks.normal((rows, HIDDEN), seed * 3 + r) for r in range(WORLD)] + outputs = hop.mnnvl_hop( + per_rank_rows[RANK], residual, gamma, EPS, comm_buffer, buffer_flags, False + ) + ref = hop.mnnvl_hop.reference(per_rank_rows, residual, gamma, EPS) + elif step == 4: + per_rank_rows = [ranks.normal((rows, HIDDEN), seed * 3 + r) for r in range(WORLD)] + outputs = torch.ops.trtllm.mnnvl_fusion_allreduce( + per_rank_rows[RANK], + gamma, + residual, + EPS, + comm_buffer, + buffer_flags, + True, + None, + int(AllReduceFusionOp.RESIDUAL_RMS_NORM), + ) + ref = hop.mnnvl_hop.reference(per_rank_rows, residual, gamma, EPS) + else: + routing = "top8" if rows > 1 else "routed_row" + per_rank = [_moe_inputs(rows, seed, r, routing) for r in range(WORLD)] + outputs = _fused(per_rank[RANK], residual, gamma) + ref = entry.mnnvl_moe_finalize_hop.reference(per_rank, residual, gamma, EPS) + issued.append((step, *outputs, *ref)) + torch.cuda.synchronize() + for step, norm, residual_out, ref_norm, ref_residual in issued: + _compare(norm, residual_out, ref_norm, ref_residual, f"interleaved step {step}") + + +def check_cuda_graph_replays() -> None: + """A captured verify-layer pair and a draft-step hop replay right with fresh routing.""" + gamma = _gamma(19) + static8 = [t.clone() for t in _moe_inputs(8, 5000, RANK, "top8")] + static1 = [t.clone() for t in _moe_inputs(1, 5001, RANK, "routed_row")] + residual8 = _residual(8, 5000).clone() + + def chain(): + norm8, res8 = _fused(static8, residual8, gamma) + norm1, res1 = _fused(static1, res8[:1].contiguous(), gamma) + return norm8, res8, norm1, res1 + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + chain() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs = chain() + for replay in range(5): + seed = 5100 + 10 * replay + per_rank8 = [_moe_inputs(8, seed, r, "top8") for r in range(WORLD)] + per_rank1 = [_moe_inputs(1, seed + 1, r, "routed_row") for r in range(WORLD)] + for dst, src in zip(static8, per_rank8[RANK]): + dst.copy_(src) + for dst, src in zip(static1, per_rank1[RANK]): + dst.copy_(src) + residual8.copy_(_residual(8, seed)) + torch.cuda.synchronize() + dist.barrier() + graph.replay() + torch.cuda.synchronize() + norm8, res8, norm1, res1 = outputs + ref = entry.mnnvl_moe_finalize_hop.reference(per_rank8, residual8, gamma, EPS) + _compare(norm8, res8, *ref, f"replay {replay} verify layer") + ref = entry.mnnvl_moe_finalize_hop.reference(per_rank1, res8[:1], gamma, EPS) + _compare(norm1, res1, *ref, f"replay {replay} draft step") + ranks.assert_same_on_every_rank(norm1, f"replay {replay} draft step") + dist.barrier() + del graph + torch.cuda.synchronize() + + +def check_the_guard_refuses_an_index_out_of_range() -> None: + """An index past the FC2 rows reads other memory as an expert row: is_valid refuses it.""" + gemm2_output, weights, index, addend = _moe_inputs(8, 6000, RANK, "top8") + index = index.clone() + index[3, 2] = gemm2_output.shape[0] + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, 8, HIDDEN) + residual = _residual(8, 6000) + try: + with validating(entry.mnnvl_moe_finalize_hop): + entry.mnnvl_moe_finalize_hop( + gemm2_output, + weights, + index, + addend, + residual, + _gamma(23), + EPS, + comm_buffer, + buffer_flags, + ) + except AssertionError as exc: + assert "must index gemm2_output's rows" in str(exc), exc + else: + raise AssertionError("is_valid admitted an index past the FC2 rows") + + +CHECKS = ( + check_certified_cells, + check_empty_slots_and_wide_rows, + check_lamport_flags_rotate, + check_interleaved_with_the_hops, + check_cuda_graph_replays, + check_the_guard_refuses_an_index_out_of_range, +) + + +def _run_one_rank() -> None: + global entry, hop, ws, validating, RANK, WORLD, MAPPING + RANK, WORLD, MAPPING = ranks.init_rank() + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.comm import ( # noqa: E501 + _mnnvl_workspace, + mnnvl_hop, + mnnvl_moe_finalize_hop, + ) + + entry, hop, ws = mnnvl_moe_finalize_hop, mnnvl_hop, _mnnvl_workspace + # Run as a script, not collected: the shared _validating module is reached the plain way. + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating as _validating_cm + + validating = _validating_cm + ranks.run_checks(CHECKS, RANK) + + +if __name__ == "__main__": + if ranks.WORKER_FLAG in sys.argv: + _run_one_rank() + ranks.spawn_ranks(Path(__file__).resolve(), DEADLINE_S) + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_ranks.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_ranks.py new file mode 100644 index 000000000000..634311451185 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_ranks.py @@ -0,0 +1,221 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Rank plumbing shared by this target's MNNVL op matrices. + +A matrix is its own launcher. Run plainly, it spawns one worker per device named in +CUDA_VISIBLE_DEVICES (``spawn_ranks``) and kills them all on the first failure or at its deadline, +because a broken Lamport exchange hangs rather than raising. Run with ``--rank-worker`` it is one +rank (``init_rank``). The workers rendezvous through torch.distributed's env:// method without +MPI, and the MNNVL workspaces are built over the TP process group, as in a non-MPI serving engine. + +The c1 world, 16 ranks on 4 nodes of 4 GPUs on one MNNVL fabric, is started by hand, one worker +per GPU, for example inside a 4-node allocation: + + export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -1) + export MASTER_PORT=29511 TRTLLM_ENABLE_PDL=1 + srun -N4 --ntasks-per-node=4 python _mnnvl_hop_op_matrix.py --rank-worker + +``init_rank`` then reads the rank, the world size and the local rank from Slurm. The checks that +need an oracle only a 1 MiB one-shot all-reduce provides skip the shapes above it. + +Only torch is imported at module scope: the launcher must carry no tensorrt_llm state. +""" + +from __future__ import annotations + +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Tuple + +import torch +import torch.distributed as dist + +WORKER_FLAG = "--rank-worker" + +# Lamport data per rank up to which the built-in MNNVL all-reduce runs one-shot, with the same +# reduction order as the hops: the shapes where it serves as a bitwise oracle. +BUILTIN_ONE_SHOT_BYTES = 1 << 20 + + +def spawn_ranks(script: Path, deadline_s: int) -> None: + """Start one worker per device in CUDA_VISIBLE_DEVICES and wait for all of them. + + The workers rendezvous through env:// on a free local port. Each gets its own session, so + that on the first nonzero exit or at the deadline every session is killed. + """ + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the devices this run owns, e.g. 0,1,2,3" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size >= 2, f"CUDA_VISIBLE_DEVICES names {world_size} device(s); need 2 or more" + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + ) + # The children rendezvous through env://, not through the launching step's PMIx: drop its + # namespace, or the first hard exit trips srun's errhandler before anything is reported. + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + print(f"[launcher] spawning {world_size} ranks, rendezvous on :{port}", flush=True) + children = [ + subprocess.Popen( + [sys.executable, str(script), WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + for r in range(world_size) + ] + + def _kill_all() -> None: + for p in children: + try: + os.killpg(p.pid, signal.SIGKILL) + except ProcessLookupError: + pass + for p in children: + p.wait() + + deadline = time.monotonic() + deadline_s + live = set(range(world_size)) + failure = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {deadline_s}s (wedged)" + break + for r in sorted(live): + code = children[r].poll() + if code is None: + continue + live.discard(r) + if code != 0: + failure = f"rank {r} exited {code}" + break + time.sleep(0.5) + if failure is not None: + _kill_all() + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + + +def init_rank() -> Tuple[int, int, Any]: + """Bind this worker's GPU, join the process group, and return ``(rank, world, mapping)``. + + ``mapping`` is the pure tensor-parallel mapping over the whole world, the c1 server's. + """ + os.environ["TLLM_DISABLE_MPI"] = "1" + if "RANK" not in os.environ and "SLURM_PROCID" in os.environ: + os.environ["RANK"] = os.environ["SLURM_PROCID"] + os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"] + os.environ["LOCAL_RANK"] = os.environ["SLURM_LOCALID"] + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ.get("LOCAL_RANK", rank)) + torch.cuda.set_device(local_rank) + # device_id keeps NCCL teardown off the implicit-device path; gloo serves the host + # agreements of the workspace construction. + dist.init_process_group( + "cuda:nccl,cpu:gloo", + init_method="env://", + rank=rank, + world_size=world, + device_id=torch.device("cuda", local_rank), + ) + from tensorrt_llm.mapping import Mapping + + gpus_per_node = min(world, torch.cuda.device_count()) + mapping = Mapping(world_size=world, rank=rank, gpus_per_node=gpus_per_node, tp_size=world) + return rank, world, mapping + + +def run_checks(checks, rank: int) -> None: + """Run every check in order; a failing rank exits hard, and the launcher kills its peers.""" + for check in checks: + try: + check() + except BaseException: + import traceback + + print(f"[rank {rank}] FAILED {check.__name__}", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + # A rank that leaves a collective early wedges its peers: exit at once. + os._exit(13) + dist.barrier() + torch.cuda.synchronize() + print(f"[rank {rank}] {len(checks)} checks passed", flush=True) + sys.stdout.flush() + # Process groups and MNNVL workspaces are torn down with the process. + os._exit(0) + + +def normal(shape, seed: int, scale: float = 1.0, dtype: torch.dtype = torch.bfloat16): + """Normal values times ``scale`` in ``dtype``, the same on every rank for one seed.""" + gen = torch.Generator(device="cuda").manual_seed(seed) + return (torch.randn(shape, generator=gen, device="cuda") * scale).to(dtype) + + +def eighths(shape, seed: int, bound: int, dtype: torch.dtype = torch.bfloat16): + """Multiples of 1/8 in [-bound / 8, bound / 8], exact in BF16 and in their sums up to 32.""" + gen = torch.Generator(device="cuda").manual_seed(seed) + raw = torch.randint(-bound, bound + 1, shape, generator=gen, device="cuda") + return (raw.float() / 8.0).to(dtype) + + +_INTEGER_OF_SIZE = {1: torch.uint8, 2: torch.int16, 4: torch.int32, 8: torch.int64} + + +def bits(t: torch.Tensor) -> torch.Tensor: + """The tensor's bits as integers of its element size.""" + return t.view(_INTEGER_OF_SIZE[t.element_size()]) + + +def assert_bitwise(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + """Gate of exactly zero, on the bits (NaN payloads and signed zeros included).""" + assert out.dtype == ref.dtype, f"{where}: dtype {out.dtype} != {ref.dtype}" + assert out.shape == ref.shape, f"{where}: shape {tuple(out.shape)} != {tuple(ref.shape)}" + differ = bits(out) != bits(ref) + if bool(differ.any()): + first = tuple(int(i) for i in differ.nonzero()[0]) + raise AssertionError( + f"{where}: {int(differ.sum())} of {out.numel()} elements differ, the first at {first}: " + f"{out[first].item()} != {ref[first].item()}" + ) + + +def assert_same_on_every_rank(t: torch.Tensor, where: str) -> None: + """Every rank of the world holds the same bits in ``t`` as rank 0.""" + gathered = [torch.empty_like(t) for _ in range(dist.get_world_size())] + dist.all_gather(gathered, t.contiguous()) + for r, other in enumerate(gathered): + assert_bitwise(other, gathered[0], f"{where}: rank {r} vs rank 0") + + +def lamport_flags(buffer_flags: torch.Tensor) -> list[int]: + """The nine Lamport flag words, as host integers.""" + torch.cuda.synchronize() + return [int(v) for v in buffer_flags.view(torch.int32).tolist()] + + +def assert_rotated(before: list[int], after: list[int], bytes_to_clear: int, where: str) -> None: + """One call later the current buffer is the dirty one, of one stage, with its bytes to clear.""" + expected = [ + (before[0] + 1) % 3, + before[0], + before[2], + 1, + bytes_to_clear, + 0, + 0, + 0, + 0, + ] + assert after == expected, f"{where}: flags {after} after {before}, expected {expected}" diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_rank_job.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_rank_job.py new file mode 100644 index 000000000000..3cc315bed18e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_rank_job.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Run one of this target's collective entries' rank jobs from pytest, and report what it did. + +The matrix of an entry, ``__op_matrix.py`` beside this file, is its own launcher: one rank +per device, its own process group, and a deadline it enforces by killing that group, which keeps a +wedged Lamport exchange from taking the calling run down with it. This only selects the devices, +starts the launcher in a fresh interpreter (it must carry no CUDA or torch.distributed state of +its own), and turns its exit code into an assertion. + +The ranks import ``tensorrt_llm`` from where this pytest process finds it, so a source checkout +and an installed wheel both work. +""" + +from __future__ import annotations + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path + +WORLD_SIZE = 4 +"""One node's ranks. The c1 world of 16 ranks on 4 nodes is started by hand (``_mnnvl_ranks``).""" + +_LAUNCHER_GRACE_S = 300 +"""Headroom over the matrix's own deadline, so its message wins the race.""" + + +def _devices() -> str: + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible: + devices = [d for d in visible.split(",") if d.strip()] + else: + import torch + + devices = [str(i) for i in range(torch.cuda.device_count())] + assert len(devices) >= WORLD_SIZE, ( + f"this matrix runs at world size {WORLD_SIZE}; only {len(devices)} device(s) are visible" + ) + return ",".join(devices[:WORLD_SIZE]) + + +def _load_launcher_constants(launcher: Path): + """Read the matrix's module-level deadline without importing it by name.""" + spec = importlib.util.spec_from_file_location(launcher.stem, launcher) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def run(entry: str) -> None: + """Run ``__op_matrix``'s launcher over ``WORLD_SIZE`` devices.""" + launcher = Path(__file__).resolve().with_name(f"_{entry}_op_matrix.py") + env = dict(os.environ, CUDA_VISIBLE_DEVICES=_devices()) + package = importlib.util.find_spec("tensorrt_llm") + assert package is not None and package.origin is not None, "tensorrt_llm is not importable" + package_parent = str(Path(package.origin).resolve().parents[1]) + env["PYTHONPATH"] = os.pathsep.join( + p for p in (package_parent, env.get("PYTHONPATH", "")) if p + ) + timeout = _load_launcher_constants(launcher).DEADLINE_S + _LAUNCHER_GRACE_S + completed = subprocess.run( + [sys.executable, str(launcher)], + env=env, + capture_output=True, + text=True, + timeout=timeout, + ) + if completed.returncode != 0: + raise AssertionError( + f"{launcher.name} exited {completed.returncode}\n" + f"--- stdout ---\n{completed.stdout}\n" + f"--- stderr ---\n{completed.stderr}" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_hop_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_hop_op_matrix.py new file mode 100644 index 000000000000..6a2ea85bc49a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_hop_op_matrix.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the mnnvl_hop entry's certification matrix. + +The matrix is ``_mnnvl_hop_op_matrix.py`` beside this file, its own 4-rank launcher on one node; +see ``_rank_job`` for why it runs as a separate job. The 16-rank c1 world is run by hand +(``_mnnvl_ranks``). +""" + +import _rank_job +import pytest + +from tensorrt_llm._utils import get_sm_version + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the glm_moe_dsa tp16 target is certified on sm_107 only" +) + + +# The job spawns 4 ranks over the visible devices; under xdist several workers would fight for +# the same GPUs. +@pytest.mark.no_xdist +def test_mnnvl_hop_op_matrix() -> None: + _rank_job.run("mnnvl_hop") diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_moe_finalize_hop_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_moe_finalize_hop_op_matrix.py new file mode 100644 index 000000000000..f3254d88db23 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_moe_finalize_hop_op_matrix.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the mnnvl_moe_finalize_hop entry's certification matrix. + +The matrix is ``_mnnvl_moe_finalize_hop_op_matrix.py`` beside this file, its own 4-rank launcher +on one node; see ``_rank_job`` for why it runs as a separate job. The 16-rank c1 world is run by +hand (``_mnnvl_ranks``). +""" + +import _rank_job +import pytest + +from tensorrt_llm._utils import get_sm_version + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the glm_moe_dsa tp16 target is certified on sm_107 only" +) + + +# The job spawns 4 ranks over the visible devices; under xdist several workers would fight for +# the same GPUs. +@pytest.mark.no_xdist +def test_mnnvl_moe_finalize_hop_op_matrix() -> None: + _rank_job.run("mnnvl_moe_finalize_hop") From d854d2afc80b69cd20fb7548733a9273792f68a5 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:11:22 -0700 Subject: [PATCH 28/53] Add the glm_moe_dsa tp16 target's MNNVL logits all-gather trtllm::glm_mnnvl_allgather gathers every rank's BF16 rows along the last dimension in rank order and widens them to FP32 in one kernel: each rank stores its shard once through the multicast address at [row][rank][column] of a Lamport buffer and copies the gathered rows out of its local copy as they land, then fills the previous call's buffer with the empty word. The result is the all-gather, concatenation and FP32 cast of the vocabulary-parallel LM head's logits bit for bit; a word of all ones (two BF16 NaNs) arrives with its lowest bit cleared. The gather runs on a LamportExchangeWorkspace: a private MNNVL workspace of the TP group whose unwritten words hold 0xFFFFFFFF, so -0.0 travels unchanged, with flags in the layout of the all-reduce workspace. The draft steps' argmax exchanges own workspaces of the same kind; construction is collective and converges failures over the group. The entry mnnvl_allgather carries the c1 cell (8 rows of the 9,680-entry vocabulary shard). Its 4-rank op matrix checks the concatenation bit for bit, special values, the flags and the refill of the previous buffer, private workspaces side by side, CUDA graph replays, late ranks and the guards. Signed-off-by: Dongfeng Yu --- .../kernels/glmMoeDsa/glmMnnvlAllgather.cu | 203 +++++++++++++++ .../kernels/glmMoeDsa/glmMnnvlAllgather.h | 62 +++++ cpp/tensorrt_llm/thop/CMakeLists.txt | 1 + cpp/tensorrt_llm/thop/glmMnnvlAllgatherOp.cpp | 100 +++++++ .../ops/comm/_mnnvl_workspace.py | 83 +++++- .../ops/comm/mnnvl_allgather.py | 93 +++++++ .../comm/_mnnvl_allgather_op_matrix.py | 243 ++++++++++++++++++ ...t_modeling_v2_mnnvl_allgather_op_matrix.py | 24 ++ 8 files changed, 806 insertions(+), 3 deletions(-) create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.cu create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.h create mode 100644 cpp/tensorrt_llm/thop/glmMnnvlAllgatherOp.cpp create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_allgather.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_allgather_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_allgather_op_matrix.py diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.cu new file mode 100644 index 000000000000..c2753b020af7 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.cu @@ -0,0 +1,203 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/common/lamportUtils.cuh" +#include "tensorrt_llm/kernels/glmMoeDsa/glmLamport.cuh" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ +namespace +{ + +constexpr int kAllgatherBlockSize = 512; +// Gathered packets each thread polls together. +constexpr int kAllgatherPacketsInFlight = 2; +// A payload word equal to the empty word is sent as this word: still a pair of BF16 NaNs. +constexpr uint32_t kAllgatherSubstituteWord = kMnnvlExchangeEmptyWord - 1U; +constexpr int kPacketBytes = static_cast(sizeof(float4)); + +union AllgatherPacket +{ + float4 packed; + uint32_t words[4]; +}; + +struct AllgatherKernelParams +{ + float4 const* input; // [numRows, shardPackets] packets of BF16 + float* output; // [numRows, nRanks * shardPackets] packets of BF16, written as FP32 + void* bufferLocal; + void* bufferMcast; + uint32_t* bufferFlags; + int numRows; + int shardPackets; + int nRanks; + int rank; +}; + +__device__ __forceinline__ bool isEmptyPacket(AllgatherPacket const& packet) +{ + bool empty = false; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + empty |= packet.words[i] == kMnnvlExchangeEmptyWord; + } + return empty; +} + +// Writes gathered packet packetIdx (eight BF16 values) to the output as FP32. A BF16 value is the upper half of the +// FP32 value it widens to: shift the bits, NaN payloads included. +__device__ __forceinline__ void storeWidened(float* output, int packetIdx, AllgatherPacket const& packet) +{ + float4* out = reinterpret_cast(output) + 2 * packetIdx; + out[0] = make_float4(__uint_as_float(packet.words[0] << 16), __uint_as_float(packet.words[0] & 0xFFFF0000U), + __uint_as_float(packet.words[1] << 16), __uint_as_float(packet.words[1] & 0xFFFF0000U)); + out[1] = make_float4(__uint_as_float(packet.words[2] << 16), __uint_as_float(packet.words[2] & 0xFFFF0000U), + __uint_as_float(packet.words[3] << 16), __uint_as_float(packet.words[3] & 0xFFFF0000U)); +} + +__global__ void __launch_bounds__(kAllgatherBlockSize) mnnvlAllgatherKernel(AllgatherKernelParams params) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + // The shard is written by the preceding kernel. Dependents read the output only after their own wait. + cudaGridDependencySynchronize(); + cudaTriggerProgrammaticLaunchCompletion(); + + GlmLamportFlags flag(params.bufferFlags, /*numStages=*/1, /*volatileRead=*/false); + auto* stageMcast = static_cast(flag.getCurLamportBuf(params.bufferMcast)); + auto const* stageLocal = static_cast(flag.getCurLamportBuf(params.bufferLocal)); + + int const tid = static_cast(blockIdx.x * blockDim.x + threadIdx.x); + int const numThreads = static_cast(gridDim.x * blockDim.x); + int const rowPackets = params.nRanks * params.shardPackets; + int const shardTotal = params.numRows * params.shardPackets; + int const gatherTotal = params.numRows * rowPackets; + + // Store the shard once through the multicast address, at [row][rank][packet] of the gathered layout. + for (int idx = tid; idx < shardTotal; idx += numThreads) + { + int const row = idx / params.shardPackets; + int const col = idx - row * params.shardPackets; + AllgatherPacket packet; + packet.packed = params.input[idx]; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + if (packet.words[i] == kMnnvlExchangeEmptyWord) + { + packet.words[i] = kAllgatherSubstituteWord; + } + } + stageMcast[row * rowPackets + params.rank * params.shardPackets + col] = packet.packed; + } + + flag.ctaArrive(); + // Empty the buffer of the previous call; it is written again two calls later. + AllgatherPacket empty; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + empty.words[i] = kMnnvlExchangeEmptyWord; + } + flag.fillDirtyLamportBuf(params.bufferLocal, empty.packed); + + // Copy the gathered rows out of the local buffer as the packets of every rank land. + for (int base = tid; base < gatherTotal; base += numThreads * kAllgatherPacketsInFlight) + { + AllgatherPacket packets[kAllgatherPacketsInFlight]; + bool pending[kAllgatherPacketsInFlight]; +#pragma unroll + for (int j = 0; j < kAllgatherPacketsInFlight; ++j) + { + int const idx = base + j * numThreads; + pending[j] = idx < gatherTotal; + if (pending[j]) + { + packets[j].packed = common::loadPackedVolatile(&stageLocal[idx]).packed; + } + } + bool anyPending = true; + while (anyPending) + { + anyPending = false; +#pragma unroll + for (int j = 0; j < kAllgatherPacketsInFlight; ++j) + { + if (!pending[j]) + { + continue; + } + int const idx = base + j * numThreads; + if (isEmptyPacket(packets[j])) + { + packets[j].packed = common::loadPackedVolatile(&stageLocal[idx]).packed; + anyPending = true; + } + else + { + storeWidened(params.output, idx, packets[j]); + pending[j] = false; + } + } + } + } + + flag.waitAndUpdate({static_cast(gatherTotal * kPacketBytes), 0, 0, 0}); +#endif +} + +} // namespace + +void mnnvlAllgather(MnnvlAllgatherParams const& params) +{ + TLLM_CHECK_WITH_INFO(common::getSMVersion() >= 90, "[GLM MNNVL all-gather] requires SM 90 or newer."); + constexpr int kEltsPerPacket = kPacketBytes / static_cast(sizeof(__nv_bfloat16)); + TLLM_CHECK_WITH_INFO(params.shardDim % kEltsPerPacket == 0, + "[GLM MNNVL all-gather] a shard row must be a multiple of %d BF16 elements, got %d", kEltsPerPacket, + params.shardDim); + if (params.numRows == 0 || params.shardDim == 0) + { + return; + } + int const shardPackets = params.shardDim / kEltsPerPacket; + int const gatherTotal = params.numRows * params.nRanks * shardPackets; + static int const kSmCount = common::getMultiProcessorCount(); + int constexpr kPacketsPerBlock = kAllgatherBlockSize * kAllgatherPacketsInFlight; + // At most one CTA per SM: every CTA is resident and stores its part of the shard, so the packets the CTAs of other + // ranks poll for are always on their way. + int const numBlocks = std::max(1, std::min((gatherTotal + kPacketsPerBlock - 1) / kPacketsPerBlock, kSmCount)); + AllgatherKernelParams const kernelParams{static_cast(params.input), params.output, + params.bufferPtrLocal, params.multicastPtr, params.bufferFlags, params.numRows, shardPackets, params.nRanks, + params.rank}; + common::launchWithPdlWhenEnabled("glmMnnvlAllgather", mnnvlAllgatherKernel, dim3(numBlocks), + dim3(kAllgatherBlockSize), 0, params.stream, kernelParams); +} + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.h new file mode 100644 index 000000000000..5abbe8bbe048 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.h @@ -0,0 +1,62 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +//! Value of every 32-bit word of an unwritten slot of an exchange workspace; a new workspace is filled with it. +constexpr uint32_t kMnnvlExchangeEmptyWord = 0xFFFFFFFFU; + +//! \brief The one-shot MNNVL all-gather of BF16 rows along the last dimension, widened to FP32, on an exchange +//! workspace: three Lamport buffers whose unwritten words hold kMnnvlExchangeEmptyWord, and their flags (the layout +//! of GlmLamportFlags). +struct MnnvlAllgatherParams +{ + int nRanks; //!< Ranks of the group + int rank; //!< This rank + int numRows; //!< Rows of the shard + int shardDim; //!< BF16 elements per row of one rank's shard, a multiple of 8 + void const* input; //!< This rank's shard [numRows, shardDim], BF16, contiguous, 16-byte aligned + float* output; //!< The gathered rows [numRows, nRanks * shardDim], FP32, contiguous, 16-byte aligned + void* bufferPtrLocal; //!< This rank's unicast address of the workspace's three Lamport buffers + void* multicastPtr; //!< Multicast address of the same buffers + uint32_t* bufferFlags; //!< Lamport flags of the workspace + cudaStream_t stream; +}; + +//! \brief Gathers every rank's shard into [row][rank][column] order, in one kernel. +//! +//! Every rank stores its shard once through the multicast address at [row][rank][column] of the gathered layout and +//! copies the gathered rows out of its local copy of the buffer as the packets land, widening BF16 to FP32 exactly; +//! it then fills the previous call's buffer with the empty word. The output equals the rank-ordered concatenation of +//! the shards bit for bit, except that a payload word equal to kMnnvlExchangeEmptyWord (a pair of BF16 NaNs with every +//! bit set, which no float conversion produces) arrives as kMnnvlExchangeEmptyWord - 1, still a pair of NaNs. The +//! kernel waits for its producer before it touches any memory, and its PDL dependents launch right after that wait. +void mnnvlAllgather(MnnvlAllgatherParams const& params); + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 9113f3678eaa..e9971fb08213 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -135,6 +135,7 @@ add_th_common_family( allreduceOp.cpp alltoallOp.cpp asyncUlyssesOp.cpp + glmMnnvlAllgatherOp.cpp glmMnnvlHopOp.cpp ulyssesPostUnscatterOp.cpp ulyssesPermuteScatterOp.cpp diff --git a/cpp/tensorrt_llm/thop/glmMnnvlAllgatherOp.cpp b/cpp/tensorrt_llm/thop/glmMnnvlAllgatherOp.cpp new file mode 100644 index 000000000000..ff992084c3d2 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmMnnvlAllgatherOp.cpp @@ -0,0 +1,100 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/mcastDevMemUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.h" +#include "tensorrt_llm/runtime/mcastDeviceMemory.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +// All-gather of BF16 rows along their last dimension, widened to FP32, over a dedicated MNNVL exchange workspace +// (comm_buffer, the unicast view of its three Lamport buffers, and buffer_flags): the rank-ordered concatenation of +// the shards (kernels/glmMoeDsa/glmMnnvlAllgather.h). +at::Tensor glmMnnvlAllgather( + at::Tensor const& input, at::Tensor& comm_buffer, at::Tensor& buffer_flags, int64_t group_size) +{ + char const* op = "glm_mnnvl_allgather"; + auto* mcastMem = tensorrt_llm::common::findMcastDevMemBuffer(comm_buffer.data_ptr()); + TORCH_CHECK(mcastMem != nullptr, op, ": comm_buffer must be obtained from a McastGPUBuffer instance"); + TORCH_CHECK(mcastMem->isMapped(), op, ": the MNNVL workspace handles are not attached"); + CHECK_TH_CUDA(input); + CHECK_CONTIGUOUS(input); + TORCH_CHECK(input.dim() >= 1, op, ": input must have at least one dimension"); + TORCH_CHECK(input.scalar_type() == at::ScalarType::BFloat16, op, ": input must be BF16, got ", input.scalar_type()); + TORCH_CHECK(static_cast(mcastMem->getWorldSize()) == group_size, op, ": group_size ", group_size, + " does not match the workspace group of ", mcastMem->getWorldSize(), " ranks"); + TORCH_CHECK(buffer_flags.element_size() == 4 && buffer_flags.numel() >= 9 && buffer_flags.is_contiguous(), op, + ": buffer_flags must hold the workspace's nine 32-bit Lamport flags"); + TORCH_CHECK(comm_buffer.device() == input.device() && buffer_flags.device() == input.device(), op, + ": the workspace must be on the device of input"); + + int64_t const shardDim = input.size(-1); + int64_t const numRows = shardDim == 0 ? 0 : input.numel() / shardDim; + TORCH_CHECK(shardDim * input.element_size() % 16 == 0, op, ": a shard row must be a multiple of 16 bytes"); + TORCH_CHECK(reinterpret_cast(input.data_ptr()) % 16 == 0, op, ": input must be 16-byte aligned"); + int64_t const gatheredBytes = numRows * group_size * shardDim * input.element_size(); + int64_t const bufferBytes = comm_buffer.numel() * comm_buffer.element_size() / 3; + TORCH_CHECK(gatheredBytes <= bufferBytes, op, ": ", gatheredBytes, " gathered bytes exceed the Lamport buffer of ", + bufferBytes, " bytes"); + + std::vector outShape(input.sizes().begin(), input.sizes().end()); + outShape.back() = group_size * shardDim; + at::Tensor out = at::empty(outShape, input.options().dtype(at::ScalarType::Float)); + if (numRows == 0) + { + return out; + } + + tensorrt_llm::kernels::glm_moe_dsa::MnnvlAllgatherParams params{}; + params.nRanks = static_cast(group_size); + params.rank = static_cast(mcastMem->getRank()); + params.numRows = static_cast(numRows); + params.shardDim = static_cast(shardDim); + params.input = input.const_data_ptr(); + params.output = out.data_ptr(); + params.bufferPtrLocal = comm_buffer.data_ptr(); + params.multicastPtr = mcastMem->getMulticastPtr(); + params.bufferFlags = reinterpret_cast(buffer_flags.data_ptr()); + params.stream = at::cuda::getCurrentCUDAStream(input.get_device()); + tensorrt_llm::kernels::glm_moe_dsa::mnnvlAllgather(params); + return out; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_mnnvl_allgather(Tensor input, Tensor(a!) comm_buffer, Tensor(b!) buffer_flags, int group_size) -> " + "Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_mnnvl_allgather", &tensorrt_llm::torch_ext::glmMnnvlAllgather); +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/_mnnvl_workspace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/_mnnvl_workspace.py index 27d5f04504b3..9a0cfe6d39ba 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/_mnnvl_workspace.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/_mnnvl_workspace.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The MNNVL Lamport workspaces of the target's hops. +"""The MNNVL Lamport workspaces of the target's hops and exchanges. The hops (``mnnvl_hop``, ``mnnvl_moe_finalize_hop``) run on the MNNVL all-reduce workspace of the engine's ``AllReduce`` (``get_or_scale_allreduce_mnnvl_workspace``): three Lamport buffers whose @@ -12,16 +12,25 @@ graph keeps the workspace it launches on alive only when it was fetched during that graph's capture. So the target fetches it for every hop it issues, eagerly and under capture, and never keeps it across forwards. + +The exchanges that are not all-reduces (``mnnvl_allgather``, the draft steps' argmax exchanges) +each own a ``LamportExchangeWorkspace``, built once and kept. """ from __future__ import annotations -from typing import Tuple +from typing import Optional, Tuple import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch.distributed.ops import get_or_scale_allreduce_mnnvl_workspace +from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, + _make_mnnvl_mcast_buffer, + _mnnvl_device_index, + _mnnvl_workspace_all_succeeded, + get_or_scale_allreduce_mnnvl_workspace, +) from tensorrt_llm.mapping import Mapping # Lamport buffers of an MNNVL all-reduce workspace. @@ -84,3 +93,71 @@ def publish_target( elements = NUM_LAMPORT_BUFFERS * workspace["buffer_size_bytes"] // torch.bfloat16.itemsize lamport = workspace["handle"].get_mc_buffer((elements,), torch.bfloat16, 0) return lamport, workspace["buffer_flags"].view(torch.int32) + + +class LamportExchangeWorkspace: + """A private MNNVL Lamport workspace of the TP group, for an exchange that is not an all-reduce. + + Three Lamport buffers whose unwritten 32-bit words hold 0xFFFFFFFF, so that every payload, + -0.0 included, travels unchanged, and their nine flags in the layout of the all-reduce + workspace's. Used by ``mnnvl_allgather`` (the verify logits) and by the draft-step kernels + that exchange each rank's argmax (index, value) packet, with the token-table row and the + hidden-half share where the next eh_proj comes from the token table (``draft_sample_glue``, + ``draft_ehproj_table``). Each of those kernels empties the buffer of its previous call with + its own threads, so every user owns its workspace, sized for its largest call. + + Construction is collective over the TP group: every rank constructs the same workspaces in the + same order, outside CUDA graph capture (the target does it once, in ``post_load_weights``), + and a failure on any rank raises on every rank. + + Args: + mapping: The tensor-parallel mapping of the group. + buffer_size_bytes: Bytes of one Lamport buffer, a positive multiple of 16. + + Attributes: + comm_buffer: This rank's unicast view of the three buffers as 1-D int32, the kernels' + ``comm_buffer`` (``exchange_buffer``) operand. + buffer_flags: The nine uint32 Lamport flags, the kernels' ``buffer_flags`` + (``exchange_flags``) operand. + """ + + def __init__(self, mapping: Mapping, buffer_size_bytes: int): + if buffer_size_bytes <= 0 or buffer_size_bytes % 16: + raise ValueError( + f"buffer_size_bytes must be a positive multiple of 16: {buffer_size_bytes}" + ) + comm = _get_mnnvl_workspace_comm(mapping) + workspace_bytes = NUM_LAMPORT_BUFFERS * buffer_size_bytes + state = None + error: Optional[Exception] = None + try: + handle = _make_mnnvl_mcast_buffer( + comm, workspace_bytes, mapping, mapping.is_multi_node() + ) + comm_buffer = handle.get_uc_buffer( + mapping.tp_rank, (workspace_bytes // 4,), torch.int32, 0 + ) + # Filled under inference mode, like the all-reduce workspace: the engine may build + # it while CUDA graph warm-up runs in inference mode. + with torch.inference_mode(): + comm_buffer.fill_(-1) + buffer_flags = torch.tensor( + [0, 2, buffer_size_bytes, 0, 0, 0, 0, 0, 0], + dtype=torch.uint32, + device=torch.device("cuda", _mnnvl_device_index(mapping)), + ) + torch.cuda.synchronize() + state = (handle, comm_buffer, buffer_flags) + except Exception as e: + error = e + # Also fences the fill: no rank stores into a peer's buffer before the peer filled it. + if not _mnnvl_workspace_all_succeeded(comm, error is None): + raise RuntimeError( + "MNNVL Lamport exchange workspace construction failed on at least one rank" + ) from error + assert state is not None + # The handle keeps the multicast allocation mapped; the communicator stays with it. + self._handle, self.comm_buffer, self.buffer_flags = state + self._comm = comm + self.mapping = mapping + self.buffer_size_bytes = buffer_size_bytes diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_allgather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_allgather.py new file mode 100644 index 000000000000..bd2144dd6fea --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_allgather.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MNNVL all-gather of BF16 rows widened to FP32, via the trtllm glm_mnnvl_allgather op.""" + +from __future__ import annotations + +from typing import Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._mnnvl_workspace import NUM_LAMPORT_BUFFERS + + +class GlmMnnvlAllgather(OpWrapper): + """Every rank's BF16 rows gathered along the last dimension in rank order, as FP32. + + ``out[..., r * shard:(r + 1) * shard]`` is rank ``r``'s ``input`` widened to FP32, the same on + every rank: the all-gather, concatenation and ``.float()`` of the vocabulary-parallel LM head's + logits, in one kernel on a private ``_mnnvl_workspace.LamportExchangeWorkspace``. Each rank + stores its shard once through the multicast address and copies the gathered rows out of its + local buffer as they land. Bit-exact, except that a pair of BF16 values whose 32 bits are all + set (NaNs no conversion produces) arrives with its lowest bit cleared, still NaNs. + + The kernel waits for its producer before it reads ``input`` and releases its PDL dependents + right after that wait: a dependent reads the output only after its own grid dependency wait. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The c1 verify step at TP16: 8 rows of the LM head's vocabulary shard. + CELLS: tuple[Cell, ...] = ( + Cell( + why="verify-step logits: 8 rows of this rank's 9,680-entry vocabulary shard " + "(154,880 over 16 ranks)", + spec=dict(rows=8, shard=9680), + ), + ) + + note = """ + Every rank calls the gather in the same order with the same shape. The gathered rows of one + call must fit one Lamport buffer of the workspace (8 rows of the 154,880-entry vocabulary are + 2.5 MB of BF16); the op refuses more. The workspace is the gather's alone: the kernel empties + the buffer of its previous call with its own threads. + """ + + def __call__( + self, + input: torch.Tensor, + comm_buffer: torch.Tensor, + buffer_flags: torch.Tensor, + group_size: int, + ) -> torch.Tensor: + return torch.ops.trtllm.glm_mnnvl_allgather(input, comm_buffer, buffer_flags, group_size) + + def reference( + self, + per_rank_inputs: Sequence[torch.Tensor], + comm_buffer: torch.Tensor | None = None, + buffer_flags: torch.Tensor | None = None, + group_size: int | None = None, + ) -> torch.Tensor: + """The rank-ordered concatenation along the last dimension, widened to FP32. + + A collective: ``per_rank_inputs`` holds every rank's ``input`` in rank order, in place of + this rank's. The other arguments are ignored. + """ + return torch.cat([rows.float() for rows in per_rank_inputs], dim=-1) + + def is_valid( + self, + input: torch.Tensor, + comm_buffer: torch.Tensor, + buffer_flags: torch.Tensor, + group_size: int, + ) -> None: + # The all-reduce workspace marks empty words with -0.0 instead of 0xFFFFFFFF: on it the + # gather takes stale words for rows that have not arrived. + assert comm_buffer.dtype == torch.int32 and comm_buffer.dim() == 1, ( + "comm_buffer must be the comm_buffer of a LamportExchangeWorkspace" + ) + # Flags of another workspace make the kernel poll and fill buffers whose state they do + # not describe: stale words are gathered, or a buffer still in use is emptied. + bytes_per_buffer = int(buffer_flags.view(torch.int32)[2].item()) + workspace_bytes = comm_buffer.numel() * comm_buffer.element_size() + assert NUM_LAMPORT_BUFFERS * bytes_per_buffer == workspace_bytes, ( + "buffer_flags must be the flags of comm_buffer's workspace" + ) + + +mnnvl_allgather = GlmMnnvlAllgather() diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_allgather_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_allgather_op_matrix.py new file mode 100644 index 000000000000..754b92766b87 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_allgather_op_matrix.py @@ -0,0 +1,243 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for the GLM-5.2 target's mnnvl_allgather entry. + +Its own launcher, like ``_mnnvl_hop_op_matrix.py`` (see ``_mnnvl_ranks``): + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _mnnvl_allgather_op_matrix.py + +The gather moves bits, so every gate is exact. Beyond the cell: special values (-0.0, infinities, +NaN payloads, the all-ones word), the rotation of the flags and the refill of the previous +buffer with the empty word, several private workspaces side by side (the gather's and the draft +steps'), CUDA graph replays, ranks arriving late, and the guards. The collected entry point is +``test_modeling_v2_mnnvl_allgather_op_matrix.py``. +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from typing import Any, List, Tuple + +import _mnnvl_ranks as ranks +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "mnnvl_allgather requires CUDA devices" + +DEADLINE_S = 1200 +SHARD = 9680 +# One Lamport buffer of the logits gather: the gathered logits of up to 54 rows of the +# 154,880-entry BF16 vocabulary. +GATHER_BUFFER_BYTES = 16 * 1024 * 1024 +EMPTY_WORD = -1 + +# Bound in the rank body: the entry and the workspace helpers import tensorrt_llm. +entry: Any = None +ws: Any = None +validating: Any = None +RANK = 0 +WORLD = 1 +MAPPING: Any = None +WORKSPACE: Any = None + + +def _shard(rows: int, seed: int, rank: int, shard: int = SHARD) -> torch.Tensor: + return ranks.normal((rows, shard), seed * 1009 + rank, scale=4.0) + + +def _gather(x: torch.Tensor, workspace: Any = None) -> torch.Tensor: + workspace = WORKSPACE if workspace is None else workspace + return entry.mnnvl_allgather(x, workspace.comm_buffer, workspace.buffer_flags, WORLD) + + +def _check(out: torch.Tensor, per_rank: List[torch.Tensor], where: str) -> None: + ranks.assert_bitwise(out, entry.mnnvl_allgather.reference(per_rank), where) + + +def check_certified_cells() -> None: + """Every cell: the concatenation bit for bit, the same on every rank.""" + for i, cell in enumerate(entry.mnnvl_allgather.CELLS): + rows, shard = cell.spec["rows"], cell.spec["shard"] + seed = 100 + 13 * i + per_rank = [_shard(rows, seed, r, shard) for r in range(WORLD)] + with validating(entry.mnnvl_allgather): + out = entry.mnnvl_allgather( + per_rank[RANK], WORKSPACE.comm_buffer, WORKSPACE.buffer_flags, WORLD + ) + assert out.shape == (rows, WORLD * shard) and out.dtype == torch.float32, cell.why + _check(out, per_rank, cell.why) + ranks.assert_same_on_every_rank(out, cell.why) + dist.barrier() + + +def check_special_values_travel_unchanged() -> None: + """-0.0, infinities, NaN payloads and subnormals arrive bit for bit; a word of all ones (two + BF16 NaNs) arrives with its lowest bit cleared, still two NaNs.""" + specials = torch.tensor( + [0x8000, 0x0000, 0x7F80, 0xFF80, 0x7FC1, 0xFFA5, 0x0001, 0x8001, 0xFFFF, 0xFFFF], + dtype=torch.int32, + ).to(torch.int16) + per_rank = [] + for r in range(WORLD): + x = _shard(8, 2000, r) + x.view(torch.int16)[r % 8, : specials.numel()] = specials.cuda() + per_rank.append(x) + expected_inputs = [] + for x in per_rank: + words = x.clone().view(torch.int32) + words[words == EMPTY_WORD] = EMPTY_WORD - 1 + expected_inputs.append(words.view(torch.bfloat16)) + out = _gather(per_rank[RANK]) + ranks.assert_bitwise(out, entry.mnnvl_allgather.reference(expected_inputs), "special values") + + +def check_flags_rotate_and_the_previous_buffer_is_emptied() -> None: + """Each call rotates the flags; the next one fills the buffer it used with the empty word.""" + previous = None + for step, rows in enumerate((8, 1, 3, 8, 2, 8, 1)): + seed = 3000 + step + per_rank = [_shard(rows, seed, r) for r in range(WORLD)] + before = ranks.lamport_flags(WORKSPACE.buffer_flags) + _check(_gather(per_rank[RANK]), per_rank, f"step {step}") + after = ranks.lamport_flags(WORKSPACE.buffer_flags) + gathered_bytes = rows * WORLD * SHARD * torch.bfloat16.itemsize + ranks.assert_rotated(before, after, gathered_bytes, f"step {step}") + if previous is not None: + index, words = previous + emptied = WORKSPACE.comm_buffer.view(3, -1)[index, :words] + assert bool((emptied == EMPTY_WORD).all()), f"step {step}: buffer {index} not emptied" + previous = (before[0], gathered_bytes // 4) + dist.barrier() + + +def check_private_workspaces_side_by_side() -> None: + """The draft steps' exchange workspaces live beside the gather's; calls alternate freely.""" + # The sizes of the draft_sample_glue workspace (8 rows of 16-byte packets per rank) and of the + # draft_ehproj_table one (the packet, a BF16 hidden row and the rank's FP32 hidden-half share). + glue = ws.LamportExchangeWorkspace(MAPPING, 8 * WORLD * 16) + table = ws.LamportExchangeWorkspace(MAPPING, 8 * WORLD * (16 + 2 * 6144 + 4 * 6144 // WORLD)) + shards = {id(glue): 8, id(table): 6144 // WORLD, id(WORKSPACE): SHARD} + for step in range(6): + workspace = (WORKSPACE, glue, table)[step % 3] + shard = shards[id(workspace)] + rows = 1 if workspace is glue else 8 + per_rank = [_shard(rows, 4000 + step, r, shard) for r in range(WORLD)] + _check(_gather(per_rank[RANK], workspace), per_rank, f"workspace {step % 3} step {step}") + dist.barrier() + + +def check_cuda_graph_replays() -> None: + """A captured gather reads its input buffer and the rotating flags at every replay.""" + rows = 8 + x_static = _shard(rows, 5000, RANK).clone() + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + _gather(x_static) + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = _gather(x_static) + for replay in range(7): + per_rank = [_shard(rows, 5100 + replay, r) for r in range(WORLD)] + x_static.copy_(per_rank[RANK]) + torch.cuda.synchronize() + dist.barrier() + graph.replay() + torch.cuda.synchronize() + _check(out, per_rank, f"replay {replay}") + dist.barrier() + del graph + torch.cuda.synchronize() + + +def check_late_ranks() -> None: + """Ranks reach back-to-back gathers at different times.""" + issued: List[Tuple[List[torch.Tensor], torch.Tensor]] = [] + for step in range(6): + rows = (8, 1, 5)[step % 3] + per_rank = [_shard(rows, 6000 + step, r) for r in range(WORLD)] + torch.cuda._sleep(((RANK * 5 + step) % WORLD) * 400_000) + issued.append((per_rank, _gather(per_rank[RANK]))) + torch.cuda.synchronize() + for step, (per_rank, out) in enumerate(issued): + _check(out, per_rank, f"late ranks step {step}") + + +def check_the_op_refuses_rows_past_the_buffer() -> None: + """Gathered rows that exceed one Lamport buffer are refused before any rank stores them.""" + rows = GATHER_BUFFER_BYTES // (WORLD * SHARD * torch.bfloat16.itemsize) + 1 + x = torch.zeros((rows, SHARD), dtype=torch.bfloat16, device="cuda") + try: + _gather(x) + except RuntimeError as exc: + assert "exceed the Lamport buffer" in str(exc), exc + else: + raise AssertionError(f"the op took {rows} rows past the buffer") + + +def check_the_guard_refuses_the_all_reduce_workspace() -> None: + """The all-reduce workspace marks empty words with -0.0: is_valid refuses it.""" + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, 8, 6144) + x = _shard(8, 7000, RANK) + try: + with validating(entry.mnnvl_allgather): + entry.mnnvl_allgather(x, comm_buffer, buffer_flags, WORLD) + except AssertionError as exc: + assert "LamportExchangeWorkspace" in str(exc), exc + else: + raise AssertionError("is_valid admitted the all-reduce workspace") + + +def check_the_guard_refuses_flags_of_another_workspace() -> None: + foreign = WORKSPACE.buffer_flags.clone() + foreign.view(torch.int32)[2] += 16 + x = _shard(8, 7100, RANK) + try: + with validating(entry.mnnvl_allgather): + entry.mnnvl_allgather(x, WORKSPACE.comm_buffer, foreign, WORLD) + except AssertionError as exc: + assert "flags of comm_buffer's workspace" in str(exc), exc + else: + raise AssertionError("is_valid admitted foreign flags") + + +CHECKS = ( + check_certified_cells, + check_special_values_travel_unchanged, + check_flags_rotate_and_the_previous_buffer_is_emptied, + check_private_workspaces_side_by_side, + check_cuda_graph_replays, + check_late_ranks, + check_the_op_refuses_rows_past_the_buffer, + check_the_guard_refuses_the_all_reduce_workspace, + check_the_guard_refuses_flags_of_another_workspace, +) + + +def _run_one_rank() -> None: + global entry, ws, validating, RANK, WORLD, MAPPING, WORKSPACE + RANK, WORLD, MAPPING = ranks.init_rank() + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.comm import ( # noqa: E501 + _mnnvl_workspace, + mnnvl_allgather, + ) + + entry, ws = mnnvl_allgather, _mnnvl_workspace + WORKSPACE = ws.LamportExchangeWorkspace(MAPPING, GATHER_BUFFER_BYTES) + # Run as a script, not collected: the shared _validating module is reached the plain way. + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating as _validating_cm + + validating = _validating_cm + ranks.run_checks(CHECKS, RANK) + + +if __name__ == "__main__": + if ranks.WORKER_FLAG in sys.argv: + _run_one_rank() + ranks.spawn_ranks(Path(__file__).resolve(), DEADLINE_S) + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_allgather_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_allgather_op_matrix.py new file mode 100644 index 000000000000..563639be3dbe --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_allgather_op_matrix.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the mnnvl_allgather entry's certification matrix. + +The matrix is ``_mnnvl_allgather_op_matrix.py`` beside this file, its own 4-rank launcher on one +node; see ``_rank_job`` for why it runs as a separate job. The 16-rank c1 world is run by hand +(``_mnnvl_ranks``). +""" + +import _rank_job +import pytest + +from tensorrt_llm._utils import get_sm_version + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the glm_moe_dsa tp16 target is certified on sm_107 only" +) + + +# The job spawns 4 ranks over the visible devices; under xdist several workers would fight for +# the same GPUs. +@pytest.mark.no_xdist +def test_mnnvl_allgather_op_matrix() -> None: + _rank_job.run("mnnvl_allgather") From 3c67cff711b4120f3c23240f18bc6971081be90c Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:23:17 -0700 Subject: [PATCH 29/53] Run the GLM-5.2 decode step's all-reduces as MNNVL hops with the published o_proj For up to 8 rows every all-reduce of the decode step is a one-shot MNNVL hop that adds the residual and applies the following RMSNorm: o_proj's epilogue publishes this rank's shard straight into the hop's workspace, the dense layers' MLP partial sums go through a plain hop, and the MoE layers' routed experts stop before their finalize step, which the post-MoE hop performs with the shared expert's output as its addend. Fused-A prefetches the attention weights that follow it into L2. A batch the kernels do not serve runs the built-in modules in the same order. Signed-off-by: Dongfeng Yu --- .../glm_5_2_nvfp4__sm_107__tp16/modeling.py | 480 ++++++++++-------- 1 file changed, 264 insertions(+), 216 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py index 013d8b32b6ff..d25e57487abd 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py @@ -49,12 +49,16 @@ from . import weights as _weights from .ops.attention.mla_kv_half_append import mla_kv_half_append from .ops.attention.sparse_mla_flash_decode import sparse_mla_flash_decode +from .ops.comm._mnnvl_workspace import hop_supported, hop_workspace, publish_target +from .ops.comm.mnnvl_hop import mnnvl_hop +from .ops.comm.mnnvl_moe_finalize_hop import mnnvl_moe_finalize_hop from .ops.embedding.fused_embedding import fused_embedding from .ops.gemm.dense_mlp_small_m import INTER_PER_RANK, dense_mlp_small_m from .ops.gemm.fused_a_gemm import MAX_TOKENS as FUSED_A_MAX_TOKENS from .ops.gemm.fused_a_gemm import OUT_FEATURES as FUSED_A_OUT_FEATURES from .ops.gemm.fused_a_gemm import fused_a_gemm from .ops.gemm.low_m_gemm import low_m_gemm +from .ops.gemm.o_proj_publish import o_proj_publish from .ops.gemm.q_b_fp8_rope import q_b_fp8_rope from .ops.moe.routed_experts import ( FUSED_ROUTER_MAX_TOKENS, @@ -100,43 +104,40 @@ class DecodeTarget: Routed to only when the batch has no context rows, which the core decides per step; nothing in this class reads the phase back. - The step is written out layer by layer for this deployment's configuration, - which ``__init__`` checks once: - - - Layers 0-2 (dense, unquantized): the attention all-reduces its own output, - the residual add and ``post_attention_layernorm`` follow, and the MLP's - all-reduce carries the residual add and the next layer's input norm. - - Layers 3-77 (MoE): both all-reduces carry the residual add and the - following norm; the experts finalize before theirs (16 ranks span nodes). - - The last all-reduce applies the model's final norm (``next_layer_layernorm``). - - Attention is MLA with the DSA indexer on the generation rows. For up to 8 - rows: q_a, the compressed KV and its rope part come from one GEMM over the - stored kv_a_proj_with_mqa weight; q_b with k_b absorbed and the q_a norm - folded writes the RoPE'd FP8 query in its epilogue, while on the aux stream - one kernel normalizes, ropes and appends the KV half to the FP8 cache; the - built-in indexer selects the 2048 rows; and the flash decode attends over - them at 4 heads with v_b folded in, ahead of o_proj. A larger batch takes the - built-in modules' path. - - The MoE block over up to 8 rows runs its shared expert as two split-K - GEMMs around SwiGLU on the aux stream, and its routed experts in one - trtllm-gen call on the main stream: the routing kernel runs the router GEMV - itself, the NVFP4 input quantizes on its own stream and joins right before - FC1, and FC2 runs on the K = 128 kernel. A larger batch takes the built-in - MoE module. - - The embedding table is whole on every rank, so the lookup is one local - kernel. The dense MLP runs as two weight-streaming kernels on this rank's - 768 intermediate columns for up to 8 rows (every c1 step: one request and - its 7 drafts); a larger batch takes the generic gated MLP. + A step of up to 8 rows (one request and its 7 draft tokens at c1) runs on + the target's kernels, layer by layer: + + - The embedding is one lookup in the whole BF16 table every rank holds. + - Attention: q_a, the compressed KV and its rope part come from one GEMM over + the stored kv_a_proj_with_mqa weight. q_b, with k_b absorbed and the q_a + norm folded, writes the RoPE'd FP8 query in its epilogue, while on the aux + stream one kernel normalizes, ropes and appends the KV half to the FP8 + cache. The built-in DSA indexer selects 2048 rows, the flash decode attends + over them at 4 heads with v_b folded in, and o_proj's epilogue publishes + this rank's shard into the all-reduce workspace. + - Every all-reduce is a one-shot MNNVL hop that adds the residual and applies + the following RMSNorm: post_attention_layernorm after attention, and after + the MLP or MoE the next layer's input_layernorm (the model's final norm + after the last layer). + - The dense layers 0-2 run this rank's 768 intermediate columns as two + weight-streaming kernels. + - The MoE layers run the shared expert as two split-K GEMMs around SwiGLU on + the aux stream, and the routed experts as one trtllm-gen call that runs the + router GEMV in its routing kernel, joins the NVFP4 input (quantized on its + own stream) right before FC1, runs FC2 on the K = 128 kernel and leaves the + finalize step to the post-MoE hop. + + A larger batch runs the built-in modules in the same order. """ def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: self.core = core config = core.model_config.pretrained_config + self.mapping = core.model_config.mapping self.hidden_size = config.hidden_size - for layer in core.layers[: core.num_hidden_layers]: + self.index_topk = config.index_topk + layers = core.layers[: core.num_hidden_layers] + for layer in layers: mla = layer.self_attn assert layer.allreduce is not None and layer.next_layer_layernorm is not None assert not mla.apply_rotary_emb and not mla.mapping.has_cp_helix() @@ -146,6 +147,15 @@ def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: ) and mla.kv_a_proj_with_mqa.weight.dtype == torch.bfloat16, ( f"layer {layer.layer_idx}: kv_a_proj_with_mqa is not the BF16 [q_a; kv_a] weight" ) + assert ( + getattr(mla.mqa, "has_fp8_kv_cache", False) + and mla.kv_cache_dtype != "fp8_ds_mla" + and mla.k_b_proj_trans.dtype == torch.bfloat16 + and mla.v_b_proj.dtype == torch.bfloat16 + and mla.v_b_proj.is_contiguous() + and mla.q_b_proj.weight.dtype == torch.bfloat16 + and mla.o_proj.weight.dtype == torch.bfloat16 + ), f"layer {layer.layer_idx}: unexpected MLA weights or KV cache layout" assert layer.input_layernorm.nvfp4_scale is None, "kv_a_proj is not NVFP4" if isinstance(layer.mlp, GatedMLP): assert ( @@ -179,29 +189,28 @@ def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: and not backend.bias and backend.w3_w1_weight.shape[-1] * 2 == config.hidden_size ), f"layer {layer.layer_idx}: routed experts are not unpadded TRTLLM-Gen NVFP4" - self.index_topk = config.index_topk + # q_b with k_b absorbed and the q_a norm weight folded, per layer: for head h, rows # [h * 576, h * 576 + 512) are k_b_trans[h] @ W_q_b[h, :qk_nope] and the next 64 rows # W_q_b[h, qk_nope:], every column times gamma_q_a; FP64, rounded to BF16 once. self._q_b_absorbed = {} - for layer in core.layers[: core.num_hidden_layers]: + # What fused-A prefetches into L2 while it waits for its input: the weights of the + # attention GEMMs that follow it. + self._l2_after_fused_a = {} + for layer in layers: mla = layer.self_attn - assert ( - getattr(mla.mqa, "has_fp8_kv_cache", False) - and mla.kv_cache_dtype != "fp8_ds_mla" - and mla.k_b_proj_trans.dtype == torch.bfloat16 - and mla.v_b_proj.dtype == torch.bfloat16 - and mla.v_b_proj.is_contiguous() - and mla.q_b_proj.weight.dtype == torch.bfloat16 - ), f"layer {layer.layer_idx}: unexpected MLA weights or KV cache layout" qk_head_dim = mla.qk_nope_head_dim + mla.qk_rope_head_dim w = mla.q_b_proj.weight.view(mla.num_heads_tp, qk_head_dim, mla.q_lora_rank).double() nope = torch.bmm(mla.k_b_proj_trans.double(), w[:, : mla.qk_nope_head_dim]) absorbed = torch.cat([nope, w[:, mla.qk_nope_head_dim :]], 1) absorbed = absorbed * mla.q_a_layernorm.weight.double() - self._q_b_absorbed[layer.layer_idx] = ( - absorbed.to(torch.bfloat16).reshape(-1, mla.q_lora_rank).contiguous() - ) + q_b = absorbed.to(torch.bfloat16).reshape(-1, mla.q_lora_rank).contiguous() + self._q_b_absorbed[layer.layer_idx] = q_b + prefetch = [q_b, mla.v_b_proj, mla.o_proj.weight] + if mla.mqa.indexer is not None: + prefetch.insert(1, mla.mqa.indexer.wq_b.weight) + self._l2_after_fused_a[layer.layer_idx] = [t for t in prefetch if t.is_contiguous()] + device = core.embed_tokens.weight.device # The attention scales the KV-half kernel writes and the flash decode reads. self._bmm1_scale = torch.empty(2, dtype=torch.float32, device=device) @@ -209,11 +218,7 @@ def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: self._unit_scale = torch.ones(1, dtype=torch.float32, device=device) self._kv_pool = None - moe_layers = [ - layer - for layer in core.layers[: core.num_hidden_layers] - if not isinstance(layer.mlp, GatedMLP) - ] + moe_layers = [layer for layer in layers if not isinstance(layer.mlp, GatedMLP)] # One exchange workspace per router call site: its launches must run in stream order. self._router_workspace = { layer.layer_idx: make_fused_router_workspace(device) for layer in moe_layers @@ -226,6 +231,24 @@ def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: self._quant_stream = torch.cuda.Stream() self._quant_fork = torch.cuda.Event() self._quant_ready = torch.cuda.Event() + self._fast_rows = {} + + def _fast(self, rows: int) -> bool: + """Whether every kernel of the step serves this many rows (decided once per shape).""" + fast = self._fast_rows.get(rows) + if fast is None: + fast = ( + rows <= FUSED_A_MAX_TOKENS + and rows <= FUSED_ROUTER_MAX_TOKENS + and bool( + torch.ops.trtllm.glm_dense_mlp_small_m_supported( + rows, self.hidden_size, INTER_PER_RANK + ) + ) + and hop_supported(rows, self.hidden_size, self.mapping.tp_size) + ) + self._fast_rows[rows] = fast + return fast def forward( self, @@ -237,106 +260,34 @@ def forward( **kwargs, ) -> torch.Tensor: core = self.core + layers = core.layers[: core.num_hidden_layers] if inputs_embeds is None: inputs_embeds = fused_embedding(input_ids, core.embed_tokens.weight) residual = inputs_embeds - hidden_states = core.layers[0].input_layernorm(inputs_embeds) - rows = hidden_states.shape[0] - small_m = torch.ops.trtllm.glm_dense_mlp_small_m_supported( - rows, self.hidden_size, INTER_PER_RANK - ) - fused_a = rows <= FUSED_A_MAX_TOKENS - for layer in core.layers[: core.num_hidden_layers]: + hidden_states = layers[0].input_layernorm(inputs_embeds) + if not self._fast(hidden_states.shape[0]): + for layer in layers: + hidden_states, residual = self._builtin_layer( + layer, position_ids, hidden_states, residual, attn_metadata + ) + return hidden_states + for layer in layers: + attn_out = self._attention(layer, position_ids, hidden_states, attn_metadata) + hidden_states, residual = self._post_attention_hop(layer, attn_out, residual) if isinstance(layer.mlp, GatedMLP): - hidden_states, residual = self._dense_layer( - layer, position_ids, hidden_states, residual, attn_metadata, fused_a, small_m + mlp = layer.mlp + # This rank's partial sum; the hop adds the 16 ranks'. + partial = dense_mlp_small_m( + hidden_states, mlp.gate_up_proj.weight, mlp.down_proj.weight ) + hidden_states, residual = self._hop(layer.next_layer_layernorm, partial, residual) else: - hidden_states, residual = self._moe_layer( - layer, position_ids, hidden_states, residual, attn_metadata, fused_a - ) + hidden_states, residual = self._moe(layer, hidden_states, residual) return hidden_states - def _experts(self, layer, hidden_states: torch.Tensor) -> torch.Tensor: - """The MoE block's output for up to 8 rows: shared expert plus routed experts.""" - mlp = layer.mlp - backend = getattr(mlp.experts, "backend", mlp.experts) - rows = hidden_states.shape[0] - - self._quant_fork.record() - with torch.cuda.stream(self._quant_stream): - self._quant_stream.wait_event(self._quant_fork) - x_fp4, x_sf = fp4_quantize( - hidden_states, backend.fc31_input_scale, backend.scaling_vector_size, False, False - ) - self._quant_ready.record() - routing = extract_routing_params(backend.routing_method) - - def routed() -> torch.Tensor: - return routed_experts( - self._router_logits[:rows], - routing.routing_bias, - x_fp4, - x_sf.flatten().view(torch.float8_e4m3fn), - backend.w3_w1_weight, - backend.w3_w1_weight_scale.view(torch.float8_e4m3fn), - None, - backend.act_alpha, - backend.act_beta, - backend.act_clamp, - backend.w2_weight, - backend.w2_weight_scale.view(torch.float8_e4m3fn), - None, - get_data_or_none(backend, "fc31_scale_c"), - get_data_or_none(backend, "fc31_alpha"), - get_data_or_none(backend, "fc2_alpha"), - backend.num_slots, - routing.top_k, - routing.n_group, - routing.topk_group, - backend.w3_w1_weight.shape[-2] // 2, - backend.slot_start, - backend.expert_size_per_partition, - routing.routed_scaling_factor, - backend.routing_method.routing_method_type, - True, - to_trtllm_gen_act_type(backend.activation_type), - None, - None, - None, - backend.max_num_tokens, - False, - self._quant_ready.cuda_event, - hidden_states, - mlp.gate.weight, - self._router_workspace[layer.layer_idx], - True, - )[0] - - def shared() -> torch.Tensor: - expert = mlp.shared_experts - gate_up = expert.gate_up_proj.weight - h = low_m_gemm(hidden_states, gate_up, hidden_states.new_empty(rows, gate_up.shape[0])) - h = swiglu(h) - return low_m_gemm( - h, expert.down_proj.weight, hidden_states.new_empty(rows, hidden_states.shape[1]) - ) - - routed_output, shared_output = maybe_execute_in_parallel( - routed, - shared, - mlp.event_dict[EventType.Main], - mlp.event_dict[EventType.MoeShared], - mlp.aux_stream, - ) - return shared_output.add_(routed_output) - - def _attention(self, layer, position_ids, hidden_states, attn_metadata, all_reduce_params, fused): - """MLA with the DSA indexer over generation rows, through o_proj.""" - if not fused: - return self._attention_builtin( - layer, position_ids, hidden_states, attn_metadata, all_reduce_params - ) + def _attention(self, layer, position_ids, hidden_states, attn_metadata) -> torch.Tensor: + """MLA with the DSA indexer over the generation rows: the attention output ahead of + o_proj, [rows, heads * v_head_dim].""" mla = layer.self_attn mqa = mla.mqa manager = attn_metadata.kv_cache_manager @@ -349,9 +300,15 @@ def _attention(self, layer, position_ids, hidden_states, attn_metadata, all_redu if self._kv_pool is None: self._kv_pool = manager.get_unique_primary_pool() - kv_a = fused_a_gemm(hidden_states, mla.kv_a_proj_with_mqa.weight, [])[:num_tokens] + kv_a = fused_a_gemm( + hidden_states, + mla.kv_a_proj_with_mqa.weight, + self._l2_after_fused_a[layer.layer_idx], + )[:num_tokens] q_a = kv_a[:, : mla.q_lora_rank] - quant_q = hidden_states.new_empty((num_tokens, mla.num_heads_tp * head_dim), dtype=torch.uint8) + quant_q = hidden_states.new_empty( + (num_tokens, mla.num_heads_tp * head_dim), dtype=torch.uint8 + ) def q_half(): q_b_fp8_rope( @@ -401,7 +358,9 @@ def kv_half(): qr, hidden_states[:num_tokens], position_ids[..., :num_tokens] ) ] - indexer._update_k_cache(indexer_intermediates[1], indexer_intermediates[2], attn_metadata) + indexer._update_k_cache( + indexer_intermediates[1], indexer_intermediates[2], attn_metadata + ) indices, _ = mqa.sparse_attn_predict( hidden_states[:num_tokens], None, @@ -430,7 +389,167 @@ def kv_half(): ) if indexer is not None: indexer.maybe_join_prev_topk_copy() - return mla.o_proj(output, all_reduce_params=all_reduce_params, layer_idx=mla.layer_idx) + return output + + def _post_attention_hop(self, layer, attn_out: torch.Tensor, residual: torch.Tensor): + """o_proj published into the all-reduce workspace, then the hop with the residual add + and post_attention_layernorm: (normed, residual).""" + rows = attn_out.shape[0] + lamport, flags = publish_target(self.mapping, rows, self.hidden_size) + comm_buffer, buffer_flags = hop_workspace(self.mapping, rows, self.hidden_size) + shard = attn_out.new_empty((rows, self.hidden_size)) + o_proj_publish( + attn_out, + layer.self_attn.o_proj.weight, + shard, + lamport, + flags, + self.mapping.tp_size, + self.mapping.tp_rank, + ) + norm = layer.post_attention_layernorm + return mnnvl_hop( + shard, residual, norm.weight, norm.variance_epsilon, comm_buffer, buffer_flags, True + ) + + def _hop(self, norm, partial: torch.Tensor, residual: torch.Tensor): + """The hop of a rank-partial sum with the residual add and `norm`: (normed, residual).""" + comm_buffer, buffer_flags = hop_workspace(self.mapping, partial.shape[0], self.hidden_size) + return mnnvl_hop( + partial, residual, norm.weight, norm.variance_epsilon, comm_buffer, buffer_flags, False + ) + + def _moe(self, layer, hidden_states: torch.Tensor, residual: torch.Tensor): + """The MoE block and the post-MoE hop that finalizes it: (normed, residual).""" + mlp = layer.mlp + backend = getattr(mlp.experts, "backend", mlp.experts) + rows = hidden_states.shape[0] + + self._quant_fork.record() + with torch.cuda.stream(self._quant_stream): + self._quant_stream.wait_event(self._quant_fork) + x_fp4, x_sf = fp4_quantize( + hidden_states, backend.fc31_input_scale, backend.scaling_vector_size, False, False + ) + self._quant_ready.record() + routing = extract_routing_params(backend.routing_method) + + def routed(): + # Unfinalized: [FC2 rows, expert weights, expanded index -> FC2 row]. + return routed_experts( + self._router_logits[:rows], + routing.routing_bias, + x_fp4, + x_sf.flatten().view(torch.float8_e4m3fn), + backend.w3_w1_weight, + backend.w3_w1_weight_scale.view(torch.float8_e4m3fn), + None, + backend.act_alpha, + backend.act_beta, + backend.act_clamp, + backend.w2_weight, + backend.w2_weight_scale.view(torch.float8_e4m3fn), + None, + get_data_or_none(backend, "fc31_scale_c"), + get_data_or_none(backend, "fc31_alpha"), + get_data_or_none(backend, "fc2_alpha"), + backend.num_slots, + routing.top_k, + routing.n_group, + routing.topk_group, + backend.w3_w1_weight.shape[-2] // 2, + backend.slot_start, + backend.expert_size_per_partition, + routing.routed_scaling_factor, + backend.routing_method.routing_method_type, + False, + to_trtllm_gen_act_type(backend.activation_type), + None, + None, + None, + backend.max_num_tokens, + False, + self._quant_ready.cuda_event, + hidden_states, + mlp.gate.weight, + self._router_workspace[layer.layer_idx], + True, + ) + + def shared() -> torch.Tensor: + expert = mlp.shared_experts + gate_up = expert.gate_up_proj.weight + h = low_m_gemm(hidden_states, gate_up, hidden_states.new_empty(rows, gate_up.shape[0])) + h = swiglu(h) + return low_m_gemm( + h, expert.down_proj.weight, hidden_states.new_empty(rows, self.hidden_size) + ) + + (fc2_output, expert_weights, permuted_rows), shared_output = maybe_execute_in_parallel( + routed, + shared, + mlp.event_dict[EventType.Main], + mlp.event_dict[EventType.MoeShared], + mlp.aux_stream, + ) + norm = layer.next_layer_layernorm + comm_buffer, buffer_flags = hop_workspace(self.mapping, rows, self.hidden_size) + return mnnvl_moe_finalize_hop( + fc2_output, + expert_weights, + permuted_rows, + shared_output, + residual, + norm.weight, + norm.variance_epsilon, + comm_buffer, + buffer_flags, + ) + + @staticmethod + def _builtin_layer(layer, position_ids, hidden_states, residual, attn_metadata): + """One decoder layer on the built-in modules, for a batch the kernels do not serve.""" + dense = isinstance(layer.mlp, GatedMLP) + hidden_states = DecodeTarget._attention_builtin( + layer, + position_ids, + hidden_states, + attn_metadata, + AllReduceParams(enable_allreduce=dense), + ) + if dense: + hidden_states, residual = layer.post_attention_layernorm(hidden_states, residual) + hidden_states = layer.mlp( + hidden_states, final_all_reduce_params=AllReduceParams(enable_allreduce=False) + ) + else: + hidden_states, residual = layer.allreduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=layer.post_attention_layernorm.weight, + eps=layer.post_attention_layernorm.variance_epsilon, + trigger_completion_at_end=False, + ), + ) + hidden_states = layer.mlp( + hidden_states, + None, + all_rank_num_tokens=attn_metadata.all_rank_num_tokens, + final_all_reduce_params=AllReduceParams(enable_allreduce=False), + do_finalize=True, + ) + return layer.allreduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=layer.next_layer_layernorm.weight, + eps=layer.next_layer_layernorm.variance_epsilon, + trigger_completion_at_end=dense, + ), + ) @staticmethod def _attention_builtin(layer, position_ids, hidden_states, attn_metadata, all_reduce_params): @@ -474,77 +593,6 @@ def _attention_builtin(layer, position_ids, hidden_states, attn_metadata, all_re indexer.maybe_join_prev_topk_copy() return mla.o_proj(output, all_reduce_params=all_reduce_params, layer_idx=mla.layer_idx) - def _dense_layer( - self, layer, position_ids, hidden_states, residual, attn_metadata, fused_a, small_m - ): - hidden_states = self._attention( - layer, - position_ids, - hidden_states, - attn_metadata, - AllReduceParams(enable_allreduce=True), - fused_a, - ) - hidden_states, residual = layer.post_attention_layernorm(hidden_states, residual) - if small_m: - # The rank partial; the all-reduce below sums it over the 16 ranks. - hidden_states = dense_mlp_small_m( - hidden_states, layer.mlp.gate_up_proj.weight, layer.mlp.down_proj.weight - ) - else: - hidden_states = layer.mlp( - hidden_states, final_all_reduce_params=AllReduceParams(enable_allreduce=False) - ) - return layer.allreduce( - hidden_states, - all_reduce_params=AllReduceParams( - fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, - residual=residual, - norm_weight=layer.next_layer_layernorm.weight, - eps=layer.next_layer_layernorm.variance_epsilon, - ), - ) - - def _moe_layer(self, layer, position_ids, hidden_states, residual, attn_metadata, fused_a): - hidden_states = self._attention( - layer, - position_ids, - hidden_states, - attn_metadata, - AllReduceParams(enable_allreduce=False), - fused_a, - ) - hidden_states, residual = layer.allreduce( - hidden_states, - all_reduce_params=AllReduceParams( - fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, - residual=residual, - norm_weight=layer.post_attention_layernorm.weight, - eps=layer.post_attention_layernorm.variance_epsilon, - trigger_completion_at_end=False, - ), - ) - if hidden_states.shape[0] <= FUSED_ROUTER_MAX_TOKENS: - hidden_states = self._experts(layer, hidden_states) - else: - hidden_states = layer.mlp( - hidden_states, - None, - all_rank_num_tokens=attn_metadata.all_rank_num_tokens, - final_all_reduce_params=AllReduceParams(enable_allreduce=False), - do_finalize=True, - ) - return layer.allreduce( - hidden_states, - all_reduce_params=AllReduceParams( - fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, - residual=residual, - norm_weight=layer.next_layer_layernorm.weight, - eps=layer.next_layer_layernorm.variance_epsilon, - trigger_completion_at_end=False, - ), - ) - class Glm52Nvfp4Sm107Tp16Core(DeepseekV3Model): """The decoder stack, with the per-step choice between the two targets.""" From 27fd45a7444d2a39b383958a1549917772aa8a1f Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:19:47 -0700 Subject: [PATCH 30/53] Add the GLM key-sharded DSA indexer top-k with an MNNVL histogram exchange trtllm::glm_dsa_hist_topk finds each verify row's exact top-2048 of one generation request whose indexer keys are split into contiguous shards over the TP group: every rank holds only its shard's FP32 logits, and the kernel exchanges 11-bit key histograms (up to three refinement levels) and the selected tokens through a private multicast workspace of three Lamport buffers, so every rank writes the same rows. Ties at the cut take the larger token. trtllm::glm_dsa_hist_shard_meta builds a rank's shard context length and block table from the request's KV length and INDEX_KEY block-table row. The kernels live in cpp/tensorrt_llm/kernels/glmMoeDsa (namespace glm_moe_dsa::hist_topk), the ops in thop/glmDsaHistTopKOp.cpp. A peer wait aborts after 120 s; PDL follows TRTLLM_ENABLE_PDL. The tp16 target gets two entries, hist_topk and hist_shard_meta, and a helper module with the shard geometry and the workspace (a McastGPUBuffer on the TP group, created once per process). The entries are certified on sm_107: hist_shard_meta on one GPU, hist_topk through a 4-rank op matrix (cells, rank agreement, buffer rotation, CUDA-graph replay); its 16-rank cells need a 4-node launch of the same matrix. Signed-off-by: Dongfeng Yu --- .../kernels/glmMoeDsa/glmDsaHistTopK.cu | 1968 +++++++++++++++++ .../kernels/glmMoeDsa/glmDsaHistTopK.h | 142 ++ cpp/tensorrt_llm/thop/CMakeLists.txt | 1 + cpp/tensorrt_llm/thop/glmDsaHistTopKOp.cpp | 120 + .../ops/dsa/_hist_topk.py | 130 ++ .../ops/dsa/hist_shard_meta.py | 132 ++ .../ops/dsa/hist_topk.py | 172 ++ .../dsa/_glm_rank_job.py | 105 + .../dsa/_hist_topk_op_matrix.py | 359 +++ .../dsa/test_modeling_v2_hist_shard_meta.py | 99 + .../test_modeling_v2_hist_topk_op_matrix.py | 27 + 11 files changed, 3255 insertions(+) create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.cu create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.h create mode 100644 cpp/tensorrt_llm/thop/glmDsaHistTopKOp.cpp create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_hist_topk.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/hist_shard_meta.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/hist_topk.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_glm_rank_job.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_hist_topk_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_shard_meta.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_topk_op_matrix.py diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.cu new file mode 100644 index 000000000000..aa3f2bee28a9 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.cu @@ -0,0 +1,1968 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.h" + +#include "tensorrt_llm/common/envUtils.h" + +#include + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::hist_topk +{ +namespace +{ + +constexpr int32_t kCtas = kCtasPerRow; +constexpr int32_t kThreads = 1024; +static_assert(kThreads / 32 == 32, "the block scans assume 32 warps"); +constexpr uint32_t kFull = 0xFFFFFFFFu; +constexpr int32_t kMaxGroups = kBins / kGroupBins; +constexpr int32_t kFineWords = kBins / 2; +static_assert(kHistWords == kFineWords + kMaxGroups, "histogram words: bin-count pairs, then group totals"); +constexpr uint32_t kEmpty = 0xFFFFFFFFu; +//! Exchanged bin counts saturate here; a group holds at most one saturated bin. +constexpr uint32_t kSat = 0xFFFFu; +static_assert(2 * int64_t{kSat} > kMaxShard, "two saturated bins would hold more keys than a shard"); +//! float4 of logits per thread (the keys are staged in shared memory). +constexpr int32_t kVecs = 3; +static_assert(kMaxShard == kCtas * kThreads * 4 * kVecs, "three float4 per thread at the largest shard"); +constexpr int32_t kKeys = 4 * kVecs; +constexpr int32_t kOutPerCta = kTopK / kCtas; +constexpr int32_t kCtasPerLaunch = kCtas * kMaxRows; +//! Boundary sub-bins of at most this many pairs are ranked directly. +constexpr int32_t kSmallRank = 64; +constexpr int64_t kHistRegionWords = int64_t{kMaxRows} * kMaxRanks * kHistWords; +constexpr int64_t kOutRegionWords = int64_t{kMaxRows} * kTopK; + +enum StateSlot : int32_t +{ + kStateCount = 0, // 64-bit: CTAs of exchange launches (each adds one once it read the count) + kStateError = 2, + kStateRecords = 4 +}; + +constexpr int32_t kRecordWords = 4; +static_assert(kStateRecords + kMaxRows * kRecordWords == kStateWords, "state layout"); + +enum Mode : int32_t +{ + kModeIdentity = 0, // no exchange for the row + kModeFull = 1, // the boundary bin is taken whole + kModeTies = 2, // the exact key: its ties are taken by count, the higher ranks first + kModeCand = 3, // the boundary pairs are exchanged and resolved on every rank + kModeNext = 4 // the boundary bin is refined with the next digit +}; + +enum DebugSlot : int32_t +{ + // %globaltimer stamps + kDbgEntry = 0, + kDbgWaited = 1, + kDbgLoaded = 2, + kDbgLocalHist = 3, + kDbgB0 = 4, + kDbgSent0 = 5, + kDbgSuffix = 6, + kDbgBucketed = 7, + kDbgCoarse = 8, + kDbgGroup = 9, + kDbgFine = 10, + kDbgCut = 11, + kDbgPeer = 12, + kDbgClassify = 13, + kDbgStaged = 14, + kDbgFinalSent = 15, + kDbgOutCopied = 16, + kDbgBndPolled = 17, + kDbgPass1 = 18, + kDbgThreshold = 19, + kDbgResolved = 20, + kDbgEnd = 21, + kDbgCleared = 22, + kDbgPreWaited = 24, + kDbgTokSent = 25, + kDbgClockEntry = 26, // clock64 (not %globaltimer) + kDbgClockEnd = 27, + // statistics + kDbgCutGroup = 32, + kDbgCutBin = 33, + kDbgOwnAbove = 34, + kDbgLevels = 35, + kDbgMode = 36, + kDbgBoundary = 37, + kDbgNeed = 38, + kDbgOutOff = 39, + kDbgBndOff = 40, + kDbgSelCount = 41, + kDbgBndCount = 42, + kDbgLocalCut = 43 +}; +static_assert(kDbgLocalCut < kDebugSlots, "debug slots"); + +// Shared memory (32-bit words). Peers never read this CTA's shared memory: the cluster's exchanges are pushes +// (st.async) that complete on the receiver's barriers. +constexpr int32_t kSmemHist = 0; // [kBins] local histogram of the current level +constexpr int32_t kSmemSuff = kSmemHist + kBins; // [kBins + 4] its suffix sums +constexpr int32_t kSmemKeys = kSmemSuff + kBins + 4; // [kVecs][kThreads] uint4 keys +constexpr int32_t kSmemStageArea = kSmemKeys + kKeys * kThreads; // [kStageWords] buckets or stage (see below) +constexpr int32_t kStageWords = kKeys * kThreads / 2; +constexpr int32_t kSmemRHist = kSmemStageArea + kStageWords; // [kBins] bucket fill, then the resolve histogram +constexpr int32_t kSmemSlots = kSmemRHist + kBins; // [kCtas][kBins / kCtas] peers' counts of our bins +constexpr int32_t kSmemCoarse = kSmemSlots + kBins; // [kMaxRanks][kMaxGroups] +constexpr int32_t kSmemFine = kSmemCoarse + kMaxRanks * kMaxGroups; // [kMaxRanks][kGroupBins] +constexpr int32_t kSmemGsum = kSmemFine + kMaxRanks * kGroupBins; // [2][kMaxGroups] group totals (all, lower ranks) +constexpr int32_t kPreRow = kGroupBins + 4; // the group's kGroupBins + 1 suffix sums, padded to 16-byte pushes +constexpr int32_t kSmemPre = kSmemGsum + 2 * kMaxGroups; // [kLevels][kCtas][kPreRow] CTAs' suffix sums +constexpr int32_t kSmemScratch = kSmemPre + kLevels * kCtas * kPreRow; // [256] scan scratch +constexpr int32_t kSmemCut = kSmemScratch + 256; // [32] the cut of the current level +constexpr int32_t kSmemOffs = kSmemCut + 32; // [8] this CTA's offsets (warp 0) +constexpr int32_t kSmemMisc = kSmemOffs + 8; // [16] +constexpr int32_t kSmemSmall = kSmemMisc + 16; // [2 * kSmallRank] +constexpr int32_t kSmemBar = kSmemSmall + 2 * kSmallRank; // [2 * kLevels] uint64 barriers +constexpr int32_t kSmemDbg = kSmemBar + 4 * kLevels; // [kDebugSlots] int64 (debug only) +constexpr int32_t kSmemWords = kSmemDbg + 2 * kDebugSlots; +constexpr size_t kSmemBytes = size_t{kSmemWords} * sizeof(uint32_t); +// The stage area holds, for a cut on the first digit, the level-0 buckets: 16-bit key indices of this CTA (the key at +// index i is keys[i / 4] component i % 4, its token base + cbase + i); for a later cut, the selected tokens. +constexpr int32_t kSmemBucket = kSmemStageArea; // [kKeys * kThreads] uint16 +constexpr int32_t kSmemStage = kSmemStageArea; // [kTopK] selected tokens +static_assert(kTopK <= kStageWords && kKeys * kThreads <= 65536, "16-bit bucket indices, the stage in the buckets"); +static_assert(kSmemKeys % 4 == 0 && kSmemStage % 4 == 0 && kSmemSlots % 4 == 0 && kSmemFine % 4 == 0 + && kSmemPre % 4 == 0 && kSmemScratch % 4 == 0 + && kSmemBar % 2 == 0 && kSmemDbg % 2 == 0, + "aligned arrays"); +// Below the 132 KB carveout: L1 keeps the register spills of 1024 threads. +static_assert(kSmemBytes <= 131 * 1024, "shared memory"); + +// kSmemCut slots. +enum CutSlot : int32_t +{ + kCutBin = 0, + kCutAboveAll = 1, // all ranks' keys above the cut bin (at this level, within the active prefix) + kCutBinAll = 2, // all ranks' keys in the cut bin + kCutAboveLow = 3, // the lower ranks' keys above the cut bin + kCutGroup = 4, // the group of the cut bin + kCutGroupAbove = 5, // all ranks' keys in the groups above it + kCutGroupAboveLow = 6, // the lower ranks' keys in the groups above it + kCutBndRankOff = 7, // the lower ranks' keys in the cut bin + kCutTiesOff = 8, // the exact key's ties the lower ranks take (when the cut is the exact key) + kCutTakeRank = 9, // the ties this rank takes + kCutRank = 16 // [kMaxRanks] every rank's count in the cut bin +}; + +// kSmemOffs slots. +enum OffsSlot : int32_t +{ + kOffsOut = 0, // this CTA's first word in the output stage + kOffsBnd = 1, // this CTA's first pair in the boundary stage + kOffsTake = 2, // the exact-key ties this CTA takes (kModeTies) + kOffsBndCta = 3, // this CTA's keys in the boundary bin + kOffsAboveCta = 4 // this CTA's keys above the boundary bin +}; + +enum MiscSlot : int32_t +{ + kMiscLevelBin = 0, // [kLevels] the cut bin of each level + kMiscSmallCount = 4, + kMiscBucketBin = 5, // the lowest bin of the level-0 buckets + kMiscRec = 6, // [3] the previous exchange launch's record of the row (levels, output words, boundary pairs) + kMiscLevelGroup = 9 // [kLevels] the group of each level's cut bin +}; + +__device__ __forceinline__ int32_t levelShift(int32_t level) +{ + return level == 0 ? 21 : (level == 1 ? 10 : 0); +} + +__device__ __forceinline__ int32_t levelBits(int32_t level) +{ + return level == 2 ? 10 : 11; +} + +__device__ __forceinline__ uint32_t orderedKey(float v) +{ + uint32_t const b = __float_as_uint(v); + uint32_t const k = (b & 0x80000000u) ? ~b : (b | 0x80000000u); + // NaN keys above +inf are clamped so that no key equals the empty word. + return k < kEmpty ? k : kEmpty - 1u; +} + +__device__ __forceinline__ uint32_t nibbleCount(uint32_t mask, int32_t k) +{ + return static_cast(__popc((mask >> (4 * k)) & 0xFu)); +} + +__device__ __forceinline__ uint4 ldVolatile4(uint32_t const* p) +{ + uint4 v; + asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(p) + : "memory"); + return v; +} + +__device__ __forceinline__ void st1(uint32_t* p, uint32_t v) +{ + asm volatile("st.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} + +__device__ __forceinline__ void st4(uint32_t* p, uint4 v) +{ + asm volatile("st.global.v4.u32 [%0], {%1, %2, %3, %4};" ::"l"(p), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w) + : "memory"); +} + +__device__ __forceinline__ uint64_t globalTimer() +{ + uint64_t t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + return t; +} + +// The debug slots of this CTA live in shared memory and are copied to HistTopKParams::debug at the exit: a stamp adds +// no global store (a later release would wait for it). +__device__ __forceinline__ int64_t* debugSlots() +{ + extern __shared__ __align__(16) uint32_t dynSmem[]; + return reinterpret_cast(dynSmem + kSmemDbg); +} + +__device__ __forceinline__ void stamp(bool dbg, int32_t slot) +{ + if (dbg && threadIdx.x == 0) + { + debugSlots()[slot] = static_cast(globalTimer()); + } +} + +__device__ __forceinline__ void putDebug(bool dbg, int32_t slot, int64_t v) +{ + if (dbg && threadIdx.x == 0) + { + debugSlots()[slot] = v; + } +} + +// Thread 0: this CTA's debug slots to HistTopKParams::debug. +__device__ __forceinline__ void flushDebug(int64_t* debug, int32_t row, int32_t c) +{ + if (debug != nullptr && threadIdx.x == 0) + { + int64_t* const g = debug + (int64_t{row} * kCtas + c) * kDebugSlots; + for (int32_t i = 0; i < kDebugSlots; ++i) + { + g[i] = debugSlots()[i]; + } + } +} + +__device__ __forceinline__ void clusterArrive() +{ + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); +} + +// An arrive without release semantics: the arriving thread does not wait for its own stores to be performed. +__device__ __forceinline__ void clusterArriveRelaxed() +{ + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); +} + +__device__ __forceinline__ void clusterWait() +{ + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); +} + +__device__ __forceinline__ void clusterSync() +{ + __syncwarp(); + clusterArrive(); + clusterWait(); +} + +__device__ __forceinline__ uint32_t smemAddr(void const* ptr) +{ + return static_cast(__cvta_generic_to_shared(ptr)); +} + +// The shared::cluster address of the same shared variable in CTA `rank` of the cluster. +__device__ __forceinline__ uint32_t peerAddr(uint32_t addr, int32_t rank) +{ + uint32_t r; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(r) : "r"(addr), "r"(rank)); + return r; +} + +__device__ __forceinline__ void barInit(uint32_t bar, uint32_t count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" ::"r"(bar), "r"(count) : "memory"); +} + +// The phase's one arrival, expecting `bytes` of asynchronous stores before it completes. +__device__ __forceinline__ void barArriveExpect(uint32_t bar, uint32_t bytes) +{ + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(bar), "r"(bytes) : "memory"); +} + +// Spins on the phase with test_wait (try_wait may suspend the thread for longer than the phase takes to complete). +__device__ __forceinline__ void barWait(uint32_t bar, uint32_t parity) +{ + asm volatile( + "{\n" + ".reg .pred done;\n" + "BAR_WAIT:\n" + "mbarrier.test_wait.parity.shared::cta.b64 done, [%0], %1;\n" + "@!done bra BAR_WAIT;\n" + "}\n" ::"r"(bar), + "r"(parity) + : "memory"); +} + +// Orders the barrier initializations before the cluster arrive that follows (peers then push to them). +__device__ __forceinline__ void fenceBarInit() +{ + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); +} + +// 16 bytes into a peer's shared memory, completing on the peer's barrier. +__device__ __forceinline__ void pushV4(uint32_t remote, uint4 v, uint32_t remoteBar) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.u32 [%0], {%1, %2, %3, %4}, [%5];" ::"r"( + remote), + "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w), "r"(remoteBar) + : "memory"); +} + +__device__ __forceinline__ void st2(uint32_t* p, uint32_t a, uint32_t b) +{ + asm volatile("st.global.v2.u32 [%0], {%1, %2};" ::"l"(p), "r"(a), "r"(b) : "memory"); +} + +__device__ __forceinline__ int64_t histOff(int32_t buf, int32_t level, int32_t row, int32_t rank) +{ + return int64_t{buf} * kBufferWords + (int64_t{level} * kMaxRows + row) * kMaxRanks * kHistWords + + int64_t{rank} * kHistWords; +} + +__device__ __forceinline__ int64_t outOff(int32_t buf, int32_t row) +{ + return int64_t{buf} * kBufferWords + kLevels * kHistRegionWords + int64_t{row} * kTopK; +} + +__device__ __forceinline__ int64_t bndOff(int32_t buf, int32_t row) +{ + return int64_t{buf} * kBufferWords + kLevels * kHistRegionWords + kOutRegionWords + + int64_t{row} * kBoundaryCap * 2; +} + +struct WaitCtx +{ + int32_t* state; + uint64_t timeoutNs; + int32_t rank; +}; + +// what: 1..3 the histograms of a level, 10 the output stage, 11 the boundary stage, 50 the previous launch's words +// before clearing; 200 + level / 300 + level: the ranks' histograms disagree; 400: the shard exceeds the logits. +__device__ __noinline__ void abortKernel(WaitCtx const& w, int32_t what) +{ + if (atomicExch(&w.state[kStateError], what) == 0) + { + printf("[glm_dsa_hist_topk] rank %d, row %d: %s (%d); aborting the kernel\n", w.rank, + static_cast(blockIdx.y), what < 200 ? "a wait for a peer exceeded the timeout" : "bad input", + what); + } + __threadfence_system(); + __trap(); +} + +struct PollClock +{ + uint64_t start = 0; + int32_t iter = 0; + + __device__ __forceinline__ void tick(WaitCtx const& w, int32_t what) + { + if ((++iter & 255) == 0) + { + uint64_t const now = globalTimer(); + if (start == 0) + { + start = now; + } + else if (now - start > w.timeoutNs) + { + abortKernel(w, what); + } + } + } +}; + +__device__ __forceinline__ bool anyEmpty(uint4 v, int32_t valid) +{ + return (valid > 0 && v.x == kEmpty) || (valid > 1 && v.y == kEmpty) || (valid > 2 && v.z == kEmpty) + || (valid > 3 && v.w == kEmpty); +} + +// The first `valid` words at p once a peer wrote them. +__device__ __forceinline__ uint4 pollVec(uint32_t const* p, int32_t valid, WaitCtx const& wc, int32_t what) +{ + uint4 v = ldVolatile4(p); + PollClock clk; + while (anyEmpty(v, valid)) + { + clk.tick(wc, what); + v = ldVolatile4(p); + } + return v; +} + +// Exclusive prefixes of N words in thread order and their block totals. One barrier; ws[0, 32 N) must not be +// rewritten before the next barrier. +template +__device__ __forceinline__ void blockExclusiveScanN( + uint32_t const (&v)[N], uint32_t* ws, uint32_t (&pre)[N], uint32_t (&tot)[N]) +{ + int32_t const lane = threadIdx.x & 31; + int32_t const warp = threadIdx.x >> 5; + uint32_t x[N]; +#pragma unroll + for (int32_t n = 0; n < N; ++n) + { + x[n] = v[n]; + } +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { +#pragma unroll + for (int32_t n = 0; n < N; ++n) + { + uint32_t const y = __shfl_up_sync(kFull, x[n], o); + if (lane >= o) + { + x[n] += y; + } + } + } + if (lane == 31) + { +#pragma unroll + for (int32_t n = 0; n < N; ++n) + { + ws[32 * n + warp] = x[n]; + } + } + __syncthreads(); + uint32_t w[N]; + uint32_t y[N]; +#pragma unroll + for (int32_t n = 0; n < N; ++n) + { + w[n] = ws[32 * n + lane]; + y[n] = w[n]; + } +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { +#pragma unroll + for (int32_t n = 0; n < N; ++n) + { + uint32_t const z = __shfl_up_sync(kFull, y[n], o); + if (lane >= o) + { + y[n] += z; + } + } + } +#pragma unroll + for (int32_t n = 0; n < N; ++n) + { + tot[n] = __shfl_sync(kFull, y[n], 31); + pre[n] = __shfl_sync(kFull, y[n] - w[n], warp) + x[n] - v[n]; + } +} + +// Exclusive suffix of v in thread order (the sum over the higher threads) and the block total. One barrier; ws[0, 32) +// as above. +__device__ __forceinline__ uint32_t blockExclusiveSuffix(uint32_t v, uint32_t* ws, uint32_t& total) +{ + int32_t const lane = threadIdx.x & 31; + int32_t const warp = threadIdx.x >> 5; + uint32_t x = v; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_down_sync(kFull, x, o); + if (lane + o < 32) + { + x += y; + } + } + if (lane == 0) + { + ws[warp] = x; + } + __syncthreads(); + uint32_t const w = ws[lane]; + uint32_t y = w; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const z = __shfl_down_sync(kFull, y, o); + if (lane + o < 32) + { + y += z; + } + } + total = __shfl_sync(kFull, y, 0); + return __shfl_sync(kFull, y - w, warp) + x - v; +} + +__device__ __forceinline__ uint32_t blockExclusiveSuffix(uint32_t v, uint32_t* ws) +{ + uint32_t total = 0; + return blockExclusiveSuffix(v, ws, total); +} + +// Position bases, in position order (float4 k major, then thread, then lane j: increasing column), of the keys +// flagged in mask a and in mask b (bit 4 k + j): the position of this thread's first flagged key of float4 k is +// ba[k] (bb[k]), the next ones follow. Also the block counts of both. One block scan (16-bit fields: a CTA flags at +// most 4 keys per thread and float4). +__device__ __forceinline__ void positionBases(uint32_t ma, uint32_t mb, uint32_t* ws, uint32_t (&ba)[kVecs], + uint32_t (&bb)[kVecs], uint32_t& na, uint32_t& nb) +{ + static_assert(kVecs == 3, "three float4 counts per mask"); + uint32_t const v[3] = {nibbleCount(ma, 0) | (nibbleCount(ma, 1) << 16), + nibbleCount(ma, 2) | (nibbleCount(mb, 0) << 16), nibbleCount(mb, 1) | (nibbleCount(mb, 2) << 16)}; + uint32_t pre[3]; + uint32_t tot[3]; + blockExclusiveScanN<3>(v, ws, pre, tot); + uint32_t const ta0 = tot[0] & 0xFFFFu; + uint32_t const ta1 = tot[0] >> 16; + uint32_t const ta2 = tot[1] & 0xFFFFu; + uint32_t const tb0 = tot[1] >> 16; + uint32_t const tb1 = tot[2] & 0xFFFFu; + uint32_t const tb2 = tot[2] >> 16; + ba[0] = pre[0] & 0xFFFFu; + ba[1] = ta0 + (pre[0] >> 16); + ba[2] = ta0 + ta1 + (pre[1] & 0xFFFFu); + bb[0] = pre[1] >> 16; + bb[1] = tb0 + (pre[2] & 0xFFFFu); + bb[2] = tb0 + tb1 + (pre[2] >> 16); + na = ta0 + ta1 + ta2; + nb = tb0 + tb1 + tb2; +} + +// The bin of a local histogram (nbins <= 2 * kThreads, two adjacent bins per thread) holding the need-th largest +// item counting from the top bin: (bin, count above, bin count) into cut[0, 3). Two barriers. +__device__ __forceinline__ void findLocalCut( + uint32_t const* h, int32_t nbins, uint32_t need, uint32_t* ws, uint32_t* cut) +{ + int32_t const t = threadIdx.x; + bool const has = 2 * t < nbins; + uint32_t const c0 = has ? h[2 * t] : 0u; + uint32_t const c1 = has ? h[2 * t + 1] : 0u; + uint32_t const above1 = blockExclusiveSuffix(c0 + c1, ws); + uint32_t const above0 = above1 + c1; + if (has) + { + if (above1 < need && need <= above1 + c1) + { + cut[0] = 2 * t + 1; + cut[1] = above1; + cut[2] = c1; + } + else if (above0 < need && need <= above0 + c0) + { + cut[0] = 2 * t; + cut[1] = above0; + cut[2] = c0; + } + } + __syncthreads(); +} + +// Every CTA of the row owns nbins / kCtas bins of the level's histogram. This CTA pushes its local counts of each +// peer's bins into that peer's slots (st.async, completing on the peer's barrier `bar`), waits for the peers' counts of +// its own bins, and stores the rank's counts of its bins through the multicast address: two per word, saturated at +// kSat, and its 32-bin group totals. The local histogram must be complete (a CTA barrier before). red: >= 16 + 128 +// words of 16-byte aligned scratch. One barrier; the multicast stores are issued by the last two warps. +__device__ __forceinline__ void pushAndSend( + uint32_t const* hist, int32_t nbins, uint32_t* slots, uint32_t bar, uint32_t* red, uint32_t* mc) +{ + int32_t const t = threadIdx.x; + int32_t const perCta = nbins / kCtas; // 256 or 128 bins: whole warps, whole groups + int32_t const groupsPerCta = perCta / kGroupBins; // 8 or 4 + int32_t const c = static_cast(blockIdx.x); + int32_t const vecs = perCta / 4; + for (int32_t i = t; i < kCtas * vecs; i += kThreads) + { + int32_t const q = i / vecs; + if (q != c) + { + int32_t const k = i - q * vecs; + pushV4(peerAddr(smemAddr(slots + c * perCta + 4 * k), q), + *reinterpret_cast(hist + q * perCta + 4 * k), peerAddr(bar, q)); + } + } + if (t < perCta) + { + barWait(bar, 0); + uint32_t s = hist[c * perCta + t]; +#pragma unroll + for (int32_t q = 0; q < kCtas; ++q) + { + s += q != c ? slots[q * perCta + t] : 0u; + } + uint32_t const group = __reduce_add_sync(kFull, s); + if ((t & 31) == 0) + { + red[t / 32] = group; + } + // Lane 8 m packs the bins 8 m .. 8 m + 7 of the warp: four words of (even, odd) bins. + uint32_t const w = min(s, kSat) | (min(__shfl_down_sync(kFull, s, 1), kSat) << 16); + uint32_t const w2 = __shfl_down_sync(kFull, w, 2); + uint32_t const w4 = __shfl_down_sync(kFull, w, 4); + uint32_t const w6 = __shfl_down_sync(kFull, w, 6); + if ((t & 7) == 0) + { + *reinterpret_cast(red + 16 + t / 2) = make_uint4(w, w2, w4, w6); + } + } + __syncthreads(); + int32_t const u = t - (kThreads - 64); + if (u >= 0 && u < perCta / 8) + { + st4(mc + (c * perCta) / 2 + 4 * u, *reinterpret_cast(red + 16 + 4 * u)); + } + else if (u >= 32 && 4 * (u - 32) < groupsPerCta) + { + st4(mc + kFineWords + c * groupsPerCta + 4 * (u - 32), *reinterpret_cast(red + 4 * (u - 32))); + } +} + +// Suffix sums of a local histogram: suff[b] = keys in bins >= b, suff[nbins] = 0 (read by the cluster's CTAs). With +// bucketBin: the highest bin b with suff[b] >= min(kTopK, suff[0]) into *bucketBin (nbins without keys). The cut of +// all ranks cannot be above it: this CTA has fewer than kTopK keys above that cut. +__device__ __forceinline__ void localSuffix( + uint32_t const* hist, int32_t nbins, uint32_t* suff, uint32_t* ws, uint32_t* bucketBin) +{ + int32_t const t = threadIdx.x; + bool const has = 2 * t < nbins; + uint32_t const c0 = has ? hist[2 * t] : 0u; + uint32_t const c1 = has ? hist[2 * t + 1] : 0u; + uint32_t total = 0; + uint32_t const above = blockExclusiveSuffix(c0 + c1, ws, total); + if (has) + { + suff[2 * t + 1] = above + c1; + suff[2 * t] = above + c1 + c0; + uint32_t const thr = min(static_cast(kTopK), total); + if (bucketBin != nullptr && thr > 0u && above < thr) + { + if (above + c1 >= thr) + { + *bucketBin = static_cast(2 * t + 1); + } + else if (above + c1 + c0 >= thr) + { + *bucketBin = static_cast(2 * t); + } + } + } + if (t == 0) + { + suff[nbins] = 0; + if (bucketBin != nullptr && total == 0u) + { + *bucketBin = static_cast(nbins); + } + } + __syncthreads(); +} + +// Clear, on this rank's copy, what the previous exchange launch wrote into buffer `buf` for this row. That launch read +// only the group totals and one group of each histogram, so every histogram word is first read written (no peer +// store of that launch can land after the clear), then cleared; the output and boundary stages were read whole. +__device__ __forceinline__ void clearPrevious(HistTopKParams const& p, int32_t buf, int32_t row, int32_t levels, + int32_t outWords, int32_t bndPairs, WaitCtx const& wc, int32_t ctaIndex, int32_t ctaCount) +{ + int32_t const c = ctaIndex; + int32_t const t = threadIdx.x; + uint4 const empty = make_uint4(kEmpty, kEmpty, kEmpty, kEmpty); + for (int32_t level = 0; level < levels; ++level) + { + // Written per rank at this level: nbins / 2 bin-count words and nbins / 32 group totals. + int32_t const nbins = 1 << levelBits(level); + int32_t const fineVecs = nbins / 8; + int32_t const rankVecs = fineVecs + nbins / (4 * kGroupBins); + uint32_t* base = p.ucWorkspace + histOff(buf, level, row, 0); + for (int32_t i = c * kThreads + t; i < p.worldSize * rankVecs; i += ctaCount * kThreads) + { + int32_t const r = i / rankVecs; + int32_t const j = i % rankVecs; + uint32_t* w = base + int64_t{r} * kHistWords + (j < fineVecs ? 4 * j : kFineWords + 4 * (j - fineVecs)); + (void) pollVec(w, 4, wc, 50); + st4(w, empty); + } + } + uint32_t* o = p.ucWorkspace + outOff(buf, row); + for (int32_t i = c * kThreads + t; i < outWords; i += ctaCount * kThreads) + { + st1(o + i, kEmpty); + } + uint32_t* b = p.ucWorkspace + bndOff(buf, row); + for (int32_t i = c * kThreads + t; i < 2 * bndPairs; i += ctaCount * kThreads) + { + st1(b + i, kEmpty); + } +} + +// The lowest set bit of a non-zero mask. +__device__ __forceinline__ int32_t lowBit(uint32_t m) +{ + return static_cast(__ffs(static_cast(m))) - 1; +} + +// Warp 0, after the final cut: every CTA's keys above it and in its bin (lane q < kCtas: CTA q's suffix sums pushed +// into pre at each level, waited for on preBar0 + 8 lv), then this CTA's offsets in the output and boundary stages +// (the rank's table from receiveCut, then the lower CTAs). Into offs[] (read after the next barrier). +__device__ __forceinline__ void cutOffsets(uint32_t const* misc, uint32_t const* pre, uint32_t preBar0, + uint32_t const* cut, int32_t level, int32_t mode, uint32_t lowTot, int32_t c, uint32_t* offs, bool dbg) +{ + int32_t const lane = static_cast(threadIdx.x & 31); + for (int32_t lv = 0; lv <= level; ++lv) + { + barWait(preBar0 + 8 * lv, 0); + } + stamp(dbg, kDbgPreWaited); + uint32_t aboveQ = 0; + uint32_t bndQ = 0; + if (lane < kCtas) + { + for (int32_t lv = 0; lv < level; ++lv) + { + uint32_t const j = misc[kMiscLevelBin + lv] - misc[kMiscLevelGroup + lv] * kGroupBins; + aboveQ += pre[(lv * kCtas + lane) * kPreRow + j + 1]; + } + uint32_t const* const pq = pre + (level * kCtas + lane) * kPreRow; + uint32_t const j = cut[kCutBin] - cut[kCutGroup] * kGroupBins; + aboveQ += pq[j + 1]; + bndQ = pq[j] - pq[j + 1]; + } + uint32_t rankOff = lowTot; + uint32_t bndRankOff = 0; + uint32_t takeQ = 0; + if (mode == kModeTies) + { + // The rank's ties go to the higher CTAs first (the larger tokens). + rankOff += cut[kCutTiesOff]; + int32_t const takeRank = static_cast(cut[kCutTakeRank]); + uint32_t xq = bndQ; +#pragma unroll + for (int32_t o = 1; o < kCtas; o <<= 1) + { + uint32_t const y = __shfl_down_sync(kFull, xq, o); + if (lane + o < kCtas) + { + xq += y; + } + } + int32_t const higherQ = static_cast(xq - bndQ); + takeQ = static_cast(max(0, min(takeRank - higherQ, static_cast(bndQ)))); + } + else + { + bndRankOff = cut[kCutBndRankOff]; + if (mode == kModeFull) + { + rankOff += bndRankOff; + takeQ = bndQ; + } + } + uint32_t const outOffC = rankOff + __reduce_add_sync(kFull, lane < c ? aboveQ + takeQ : 0u); + uint32_t const bndOffC = bndRankOff + __reduce_add_sync(kFull, lane < c ? bndQ : 0u); + uint32_t const takeC = __shfl_sync(kFull, takeQ, c); + uint32_t const bndC = __shfl_sync(kFull, bndQ, c); + uint32_t const aboveC = __shfl_sync(kFull, aboveQ, c); + if (lane == 0) + { + offs[kOffsOut] = outOffC; + offs[kOffsBnd] = bndOffC; + offs[kOffsTake] = takeC; + offs[kOffsBndCta] = bndC; + offs[kOffsAboveCta] = aboveC; + } +} + +// Store n words from shared memory through the multicast address: 16-byte stores where dst is aligned. +__device__ __forceinline__ void storeWords(uint32_t* dst, uint32_t const* src, int32_t n) +{ + int32_t const t = threadIdx.x; + int32_t const head + = min(n, static_cast((4u - static_cast((reinterpret_cast(dst) >> 2) & 3u)) & 3u)); + int32_t const vecs = (n - head) / 4; + int32_t const tail = n - head - 4 * vecs; + if (t < head) + { + st1(dst + t, src[t]); + } + for (int32_t i = t; i < vecs; i += kThreads) + { + int32_t const j = head + 4 * i; + st4(dst + j, make_uint4(src[j], src[j + 1], src[j + 2], src[j + 3])); + } + if (t >= 64 && t < 64 + tail) + { + int32_t const j = head + 4 * vecs + (t - 64); + st1(dst + j, src[j]); + } +} + +// Every CTA of the row: the cut of one level from all ranks' histograms, read through the group totals (4 KB) and the +// crossing group's bins (1 KB), with every rank's count in the cut bin and the rank's offsets that follow from them. +// Warps 2..8 meanwhile push this CTA's suffix sums over the crossing group's bins and the bin above them into row c +// of every CTA's pre rows of this level (completing on their barriers preBar0 + 8 level); when this is the final cut, +// warp 0 then computes this CTA's offsets (cutOffsets). gsum must be zero; lowTotPrev: the lower ranks' keys above +// the cuts of the earlier levels. Results in cut[] and offs[] (every thread can read them after the last barrier). +__device__ void receiveCut(HistTopKParams const& p, uint32_t const* ucLevel, uint32_t const* suffLevel, int32_t nbins, + int32_t level, uint32_t need, uint32_t lowTotPrev, uint32_t* coarse, uint32_t* fine, uint32_t* gsum, + uint32_t const* misc, uint32_t* pre, uint32_t preBar0, uint32_t* cut, uint32_t* offs, WaitCtx const& wc, + int32_t what, bool dbg) +{ + uint32_t* const preL = pre + level * kCtas * kPreRow; + uint32_t const preBar = preBar0 + 8 * level; + int32_t const t = threadIdx.x; + int32_t const world = p.worldSize; + int32_t const rank = p.rank; + int32_t const ngroups = nbins / kGroupBins; // 64 or 32 + int32_t const vecsPerRank = ngroups / 4; // 16 or 8 + uint32_t* const gAllSum = gsum; // [kMaxGroups] all ranks + uint32_t* const gLowSum = gsum + kMaxGroups; // [kMaxGroups] the lower ranks + if (t == 0) + { + cut[kCutGroup] = kEmpty; + cut[kCutBin] = kEmpty; + } + // 1. The group totals of every rank, summed per group over all ranks and over the lower ranks. + if (t < kMaxRanks * vecsPerRank) + { + int32_t const r = t / vecsPerRank; + int32_t const vi = t % vecsPerRank; + uint4 v = make_uint4(0, 0, 0, 0); + if (r < world) + { + v = pollVec(ucLevel + int64_t{r} * kHistWords + kFineWords + 4 * vi, 4, wc, what); + uint32_t const vv[4] = {v.x, v.y, v.z, v.w}; +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + atomicAdd(&gAllSum[4 * vi + j], vv[j]); + if (r < rank) + { + atomicAdd(&gLowSum[4 * vi + j], vv[j]); + } + } + } + coarse[r * kMaxGroups + 4 * vi] = v.x; + coarse[r * kMaxGroups + 4 * vi + 1] = v.y; + coarse[r * kMaxGroups + 4 * vi + 2] = v.z; + coarse[r * kMaxGroups + 4 * vi + 3] = v.w; + } + __syncthreads(); + stamp(dbg, kDbgCoarse); + // 2. The crossing group (warp 0): lane l holds groups 2 l and 2 l + 1. + if (t < 32) + { + uint32_t gAll[2]; + uint32_t gLow[2]; +#pragma unroll + for (int32_t k = 0; k < 2; ++k) + { + int32_t const g = 2 * t + k; + gAll[k] = g < ngroups ? gAllSum[g] : 0u; + gLow[k] = g < ngroups ? gLowSum[g] : 0u; + } + // Exclusive suffixes over the lanes of the per-lane totals. + uint32_t xa = gAll[0] + gAll[1]; + uint32_t xl = gLow[0] + gLow[1]; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const ya = __shfl_down_sync(kFull, xa, o); + uint32_t const yl = __shfl_down_sync(kFull, xl, o); + if (t + o < 32) + { + xa += ya; + xl += yl; + } + } + uint32_t const aAbove1 = xa - gAll[0] - gAll[1]; + uint32_t const lAbove1 = xl - gLow[0] - gLow[1]; + int32_t which = -1; + if (aAbove1 < need && need <= aAbove1 + gAll[1]) + { + which = 1; + } + else if (aAbove1 + gAll[1] < need && need <= aAbove1 + gAll[1] + gAll[0]) + { + which = 0; + } + if (which >= 0) + { + cut[kCutGroup] = 2 * t + which; + cut[kCutGroupAbove] = which == 1 ? aAbove1 : aAbove1 + gAll[1]; + cut[kCutGroupAboveLow] = which == 1 ? lAbove1 : lAbove1 + gLow[1]; + } + } + __syncthreads(); + int32_t const g = static_cast(cut[kCutGroup]); + if (g < 0 || g >= ngroups) + { + // No group holds the need-th key: the ranks' inputs disagree, nothing exact can follow. + abortKernel(wc, 200 + what); + } + uint32_t const needIn = need - cut[kCutGroupAbove]; + stamp(dbg, kDbgGroup); + // Warps 2..8 (idle in steps 3 and 4): this CTA's suffix sums over the group's bins and the bin above them + // (vectors 0..8 of four, the last three words padding) to every CTA of the row, one destination per warp: the 9 + // lanes' 16-byte pushes go to consecutive addresses of one peer, and the warps push at once. + if (t >= 64 && t < 64 + 32 * (kCtas - 1)) + { + int32_t const c = static_cast(blockIdx.x); + int32_t const lane = t & 31; + int32_t const qi = (t >> 5) - 2; + int32_t const q = qi < c ? qi : qi + 1; + if (lane < kPreRow / 4) + { + uint32_t const* const src = suffLevel + g * kGroupBins + 4 * lane; + uint4 const val = make_uint4(src[0], src[1], src[2], src[3]); + uint4* const slot = reinterpret_cast(preL + c * kPreRow) + lane; + if (qi == 0) + { + *slot = val; // this CTA's own row + } + pushV4(peerAddr(smemAddr(slot), q), val, peerAddr(preBar, q)); + } + } + // 3. The group's 32 bins of every rank (threads 4 r .. 4 r + 3: 8 bins each), a saturated bin restored from the + // group total. + if (t < kMaxRanks * 4) + { + int32_t const r = t / 4; + int32_t const vi = t % 4; + uint4 v = make_uint4(0, 0, 0, 0); + if (r < world) + { + v = pollVec(ucLevel + int64_t{r} * kHistWords + (kGroupBins / 2) * g + 4 * vi, 4, wc, what); + } + uint32_t const ww[4] = {v.x, v.y, v.z, v.w}; + uint32_t x[8]; + uint32_t others = 0; + bool sat = false; +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + x[2 * j] = ww[j] & 0xFFFFu; + x[2 * j + 1] = ww[j] >> 16; + } +#pragma unroll + for (int32_t j = 0; j < 8; ++j) + { + if (x[j] == kSat) + { + sat = true; + } + else + { + others += x[j]; + } + } + others += __shfl_xor_sync(kFull, others, 1); + others += __shfl_xor_sync(kFull, others, 2); + if (sat) + { + uint32_t const total = coarse[r * kMaxGroups + g]; +#pragma unroll + for (int32_t j = 0; j < 8; ++j) + { + if (x[j] == kSat) + { + x[j] = total - others; + } + } + } + *reinterpret_cast(fine + r * kGroupBins + 8 * vi) = make_uint4(x[0], x[1], x[2], x[3]); + *reinterpret_cast(fine + r * kGroupBins + 8 * vi + 4) = make_uint4(x[4], x[5], x[6], x[7]); + // Warps 0 and 1 only: warps 2..8 may still push. + asm volatile("bar.sync 1, 64;" ::: "memory"); + } + stamp(dbg, kDbgFine); + // 4. The cut bin (warp 0: lane = bin of the group), every rank's count in it and the rank's offsets. + if (t < 32) + { + uint32_t all = 0; + uint32_t low = 0; +#pragma unroll + for (int32_t r = 0; r < kMaxRanks; ++r) + { + uint32_t const x = fine[r * kGroupBins + t]; + all += x; + low += r < rank ? x : 0u; + } + // Exclusive suffix over the 32 bins of the group. + uint32_t xa = all; + uint32_t xl = low; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const ya = __shfl_down_sync(kFull, xa, o); + uint32_t const yl = __shfl_down_sync(kFull, xl, o); + if (t + o < 32) + { + xa += ya; + xl += yl; + } + } + uint32_t const aAbove = xa - all; + uint32_t const lAbove = xl - low; + uint32_t const holder = __ballot_sync(kFull, aAbove < needIn && needIn <= aAbove + all); + if (holder != 0u) + { + int32_t const w = lowBit(holder); + uint32_t const aAboveW = __shfl_sync(kFull, aAbove, w); + uint32_t const allW = __shfl_sync(kFull, all, w); + uint32_t const lAboveW = __shfl_sync(kFull, lAbove, w); + uint32_t const bc = t < kMaxRanks ? fine[t * kGroupBins + w] : 0u; + if (t < kMaxRanks) + { + cut[kCutRank + t] = bc; + } + uint32_t const bndRankOff = __reduce_add_sync(kFull, t < rank ? bc : 0u); + // When the bin is the exact key: rank r takes what the higher ranks leave of the bin's need. + uint32_t const needBin = needIn - aAboveW; + uint32_t xs = bc; +#pragma unroll + for (int32_t o = 1; o < kMaxRanks; o <<= 1) + { + uint32_t const y = __shfl_down_sync(kFull, xs, o); + if (t + o < kMaxRanks) + { + xs += y; + } + } + int32_t const higher = static_cast(xs - bc); + int32_t const take = max(0, min(static_cast(needBin) - higher, static_cast(bc))); + uint32_t const tiesOff = __reduce_add_sync(kFull, t < rank ? static_cast(take) : 0u); + uint32_t const takeRank = __shfl_sync(kFull, static_cast(take), rank); + uint32_t const aboveAll = cut[kCutGroupAbove] + aAboveW; + uint32_t const aboveLow = cut[kCutGroupAboveLow] + lAboveW; + if (t == 0) + { + cut[kCutBin] = static_cast(g * kGroupBins + w); + cut[kCutAboveAll] = aboveAll; + cut[kCutBinAll] = allW; + cut[kCutAboveLow] = aboveLow; + cut[kCutBndRankOff] = bndRankOff; + cut[kCutTiesOff] = tiesOff; + cut[kCutTakeRank] = takeRank; + } + __syncwarp(); + stamp(dbg, kDbgCut); + // When this is the final cut (as the main loop decides it), this CTA's offsets at once. + uint32_t const needAfter = need - aboveAll; + int32_t const mode = needAfter == allW ? kModeFull + : (level == kLevels - 1 ? kModeTies + : (allW <= static_cast(kBoundaryCap) ? kModeCand : kModeNext)); + if (mode != kModeNext) + { + cutOffsets(misc, pre, preBar0, cut, level, mode, lowTotPrev + aboveLow, + static_cast(blockIdx.x), offs, dbg); + } + } + } + __syncthreads(); + stamp(dbg, kDbgPeer); +} + +// Store the tokens of n bucket entries (key indices) through the multicast address: 16-byte stores where dst is +// aligned. +__device__ __forceinline__ void storeBucketTokens(uint32_t* dst, uint16_t const* idx, uint32_t tokenBase, int32_t n) +{ + int32_t const t = threadIdx.x; + int32_t const head + = min(n, static_cast((4u - static_cast((reinterpret_cast(dst) >> 2) & 3u)) & 3u)); + int32_t const vecs = (n - head) / 4; + int32_t const tail = n - head - 4 * vecs; + if (t < head) + { + st1(dst + t, tokenBase + idx[t]); + } + for (int32_t i = t; i < vecs; i += kThreads) + { + int32_t const j = head + 4 * i; + st4(dst + j, + make_uint4(tokenBase + idx[j], tokenBase + idx[j + 1], tokenBase + idx[j + 2], tokenBase + idx[j + 3])); + } + if (t >= 64 && t < 64 + tail) + { + int32_t const j = head + 4 * vecs + (t - 64); + st1(dst + j, tokenBase + idx[j]); + } +} + +// Store n bucket entries as (key, token) pairs through the multicast address: 16-byte stores (two pairs) where dst +// is aligned (dst is 8-byte aligned). From the second half of the CTA, beside storeBucketTokens. +__device__ __forceinline__ void storeBucketPairs( + uint32_t* dst, uint16_t const* idx, uint32_t const* keys, uint32_t tokenBase, int32_t n) +{ + int32_t const t = (static_cast(threadIdx.x) + kThreads / 2) % kThreads; + int32_t const head = min(n, ((reinterpret_cast(dst) >> 3) & 1u) != 0u ? 1 : 0); + int32_t const vecs = (n - head) / 2; + if (t < head) + { + st1(dst, keys[idx[0]]); + st1(dst + 1, tokenBase + idx[0]); + } + for (int32_t i = t; i < vecs; i += kThreads) + { + int32_t const j = head + 2 * i; + uint32_t const a = idx[j]; + uint32_t const b = idx[j + 1]; + st4(dst + 2 * j, make_uint4(keys[a], tokenBase + a, keys[b], tokenBase + b)); + } + if (t == 64 && head + 2 * vecs < n) + { + uint32_t const a = idx[n - 1]; + st1(dst + 2 * (n - 1), keys[a]); + st1(dst + 2 * (n - 1) + 1, tokenBase + a); + } +} + +// The resolver: among all ranks' boundary pairs (thread t holds pairs 2 t, 2 t + 1, 2 kThreads + 2 t and +// 2 kThreads + 2 t + 1 in pk/pt, valid as flagged in pv; every key shares its bits from `shift` up with keyTop), +// select the `need` largest by (key, token) and write their tokens, in pair order, to outRow[base, base + need). +__device__ void resolveBoundary(uint32_t const (&pk)[4], uint32_t const (&pt)[4], uint32_t pv, uint32_t need, + int32_t shift, uint32_t keyTop, uint32_t* rh, uint32_t* ws, uint32_t* cut, uint32_t* small, uint32_t* smallCount, + int32_t* outRow, int32_t base, bool dbg) +{ + int32_t const t = threadIdx.x; + // The next 11 key bits (rh and smallCount are zero: cleared at the kernel start). + int32_t const w1 = shift > 11 ? 11 : shift; + int32_t const sh1 = shift - w1; + uint32_t const mask1 = (1u << w1) - 1u; +#pragma unroll + for (int32_t i = 0; i < 4; ++i) + { + if ((pv >> i) & 1u) + { + atomicAdd(&rh[(pk[i] >> sh1) & mask1], 1u); + } + } + __syncthreads(); + findLocalCut(rh, 1 << w1, need, ws, cut); + uint32_t const d1 = cut[0]; + uint32_t const need1 = need - cut[1]; + uint32_t const cnt1 = cut[2]; + stamp(dbg, kDbgPass1); + uint32_t tKey = 0; + uint32_t tTok = 0; + if (cnt1 <= static_cast(kSmallRank)) + { + // Few pairs in that sub-bin: rank them by (key, token) directly. +#pragma unroll + for (int32_t i = 0; i < 4; ++i) + { + if (((pv >> i) & 1u) && ((pk[i] >> sh1) & mask1) == d1) + { + uint32_t const slot = atomicAdd(smallCount, 1u); + small[slot] = pk[i]; + small[kSmallRank + slot] = pt[i]; + } + } + __syncthreads(); + if (t < static_cast(cnt1)) + { + uint32_t const k = small[t]; + uint32_t const tok = small[kSmallRank + t]; + uint32_t higher = 0; + for (uint32_t j = 0; j < cnt1; ++j) + { + uint32_t const kj = small[j]; + uint32_t const tj = small[kSmallRank + j]; + higher += (kj > k || (kj == k && tj > tok)) ? 1u : 0u; + } + if (higher == need1 - 1) + { + cut[0] = k; + cut[1] = tok; + } + } + __syncthreads(); + tKey = cut[0]; + tTok = cut[1]; + } + else + { + // The remaining key bits, then the token bits of the threshold key's ties (the larger tokens first). + uint32_t needR = need1; + uint32_t prefix = d1; + uint32_t tie = cnt1; + if (sh1 > 0) + { + __syncthreads(); + for (int32_t i = t; i < (1 << sh1); i += kThreads) + { + rh[i] = 0; + } + __syncthreads(); + uint32_t const mask2 = (1u << sh1) - 1u; +#pragma unroll + for (int32_t i = 0; i < 4; ++i) + { + if (((pv >> i) & 1u) && ((pk[i] >> sh1) & mask1) == d1) + { + atomicAdd(&rh[pk[i] & mask2], 1u); + } + } + __syncthreads(); + findLocalCut(rh, 1 << sh1, needR, ws, cut); + prefix = (d1 << sh1) | cut[0]; + needR -= cut[1]; + tie = cut[2]; + __syncthreads(); + } + uint32_t const kLow = prefix; + uint32_t const lowMask = (1u << shift) - 1u; + uint32_t tokPrefix = 0; + if (needR != tie) + { + int32_t rem = 21; + while (rem > 0) + { + int32_t const width = rem > 11 ? 11 : rem; + int32_t const sh = rem - width; + uint32_t const matchMask = (1u << (21 - rem)) - 1u; + for (int32_t i = t; i < (1 << width); i += kThreads) + { + rh[i] = 0; + } + __syncthreads(); +#pragma unroll + for (int32_t i = 0; i < 4; ++i) + { + if (((pv >> i) & 1u) && (pk[i] & lowMask) == kLow && ((pt[i] >> rem) & matchMask) == tokPrefix) + { + atomicAdd(&rh[(pt[i] >> sh) & ((1u << width) - 1u)], 1u); + } + } + __syncthreads(); + findLocalCut(rh, 1 << width, needR, ws, cut); + tokPrefix = (tokPrefix << width) | cut[0]; + needR -= cut[1]; + rem = sh; + __syncthreads(); + } + } + tKey = keyTop | kLow; + tTok = tokPrefix; + } + stamp(dbg, kDbgThreshold); + // The pairs at or above (tKey, tTok), in pair order. + uint32_t win = 0; +#pragma unroll + for (int32_t i = 0; i < 4; ++i) + { + if (((pv >> i) & 1u) && (pk[i] > tKey || (pk[i] == tKey && pt[i] >= tTok))) + { + win |= 1u << i; + } + } + uint32_t const v[1] = {static_cast(__popc(win & 3u)) | (static_cast(__popc(win >> 2)) << 16)}; + uint32_t pre[1]; + uint32_t tot[1]; + blockExclusiveScanN<1>(v, ws, pre, tot); + uint32_t p0 = pre[0] & 0xFFFFu; + uint32_t p1 = (tot[0] & 0xFFFFu) + (pre[0] >> 16); +#pragma unroll + for (int32_t i = 0; i < 4; ++i) + { + if ((win >> i) & 1u) + { + uint32_t& pos = i < 2 ? p0 : p1; + outRow[base + static_cast(pos)] = static_cast(pt[i]); + ++pos; + } + } +} + +__global__ void __launch_bounds__(kThreads, 1) dsaHistTopKKernel(HistTopKParams p) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + extern __shared__ __align__(16) uint32_t smem[]; + uint32_t* const hist = smem + kSmemHist; + uint32_t* const suff = smem + kSmemSuff; + uint4* const keys = reinterpret_cast(smem + kSmemKeys); + uint32_t* const stageTok = smem + kSmemStage; + uint16_t* const bucket = reinterpret_cast(smem + kSmemBucket); + uint32_t* const rhist = smem + kSmemRHist; + uint32_t* const cursor = rhist; // the level-0 bucket fill; its bins are cleared again for the resolver + uint32_t* const slots = smem + kSmemSlots; + uint32_t* const coarse = smem + kSmemCoarse; + uint32_t* const fine = smem + kSmemFine; + uint32_t* const gsum = smem + kSmemGsum; + uint32_t* const pre = smem + kSmemPre; + uint32_t* const ws = smem + kSmemScratch; + uint32_t* const cut = smem + kSmemCut; + uint32_t* const offs = smem + kSmemOffs; + uint32_t* const misc = smem + kSmemMisc; + uint32_t* const small = smem + kSmemSmall; + // Barriers: [kLevels] the peers' histogram counts, then [kLevels] the peers' suffix sums at the cut. + uint32_t const histBar0 = smemAddr(smem + kSmemBar); + uint32_t const preBar0 = histBar0 + 8 * kLevels; + + int32_t const c = static_cast(blockIdx.x); + int32_t const row = static_cast(blockIdx.y); + int32_t const t = threadIdx.x; + bool const dbg = p.debug != nullptr; + if (dbg && t == 0) + { + for (int32_t i = 0; i < kDebugSlots; ++i) + { + debugSlots()[i] = 0; + } + } + stamp(dbg, kDbgEntry); + putDebug(dbg, kDbgClockEntry, static_cast(clock64())); + + // The request's KV length is written before the logits of this call (the forward's metadata). + int32_t const kvLen = p.kvLens[0]; + int32_t const world = p.worldSize; + int32_t const rank = p.rank; + int32_t const perRank = (max(kvLen, 0) + world - 1) / world; + int32_t const S = (perRank + kShardAlign - 1) / kShardAlign * kShardAlign; + int32_t const base = rank * S; + int32_t const n = max(0, min(kvLen - base, S)); + int32_t const rowWin = kvLen - p.nextN + row + 1; + bool const rowActive = row < p.numRows; + bool const exchangeLaunch = kvLen - p.nextN + p.numRows > kTopK; + bool const rowExchange = rowActive && rowWin > kTopK; + int32_t const v = max(0, min(rowWin - base, n)); + int32_t const chunk = S / kCtas; + int32_t const cbase = c * chunk; + + for (int32_t i = t; i < kBins; i += kThreads) + { + hist[i] = 0; + rhist[i] = 0; + } + if (t == 0) + { + misc[kMiscSmallCount] = 0; + if (rowExchange) + { + // One arrival per phase (this thread, now), completing when the peers' bytes arrived. + for (int32_t lv = 0; lv < kLevels; ++lv) + { + barInit(histBar0 + 8 * lv, 1); + barInit(preBar0 + 8 * lv, 1); + barArriveExpect(histBar0 + 8 * lv, (kCtas - 1) * (1 << levelBits(lv)) / kCtas * 4); + barArriveExpect(preBar0 + 8 * lv, (kCtas - 1) * kPreRow * 4); + } + fenceBarInit(); + } + } + __syncthreads(); // the zeroed histograms before any thread counts into them, the barriers before their waits + if (rowExchange) + { + // Every CTA's barriers are initialized before any peer pushes to them (waited for after the logits load). + __syncwarp(); + clusterArriveRelaxed(); + } + + asm volatile("griddepcontrol.wait;" ::: "memory"); + stamp(dbg, kDbgWaited); + + int32_t* const outRow = p.out + int64_t{row} * kTopK; + if (rowActive && !rowExchange && t < kOutPerCta) + { + int32_t const pos = c * kOutPerCta + t; + outRow[pos] = pos < rowWin ? pos : -1; + } + if (!exchangeLaunch) + { + flushDebug(p.debug, row, c); + return; + } + WaitCtx const wc{p.state, p.timeoutNs, rank}; + if (S > p.localCapacity || S > kMaxShard) + { + abortKernel(wc, 400); + } + + // This CTA's chunk of the row's shard logits: keys into shared memory (each thread reads back only its own). + float const* rowLogits = p.logits + int64_t{row} * p.logitsStride + cbase; + float4 f[kVecs]; +#pragma unroll + for (int32_t k = 0; k < kVecs; ++k) + { + int32_t const q = t + k * kThreads; + f[k] = make_float4(0.F, 0.F, 0.F, 0.F); + if (rowExchange && 4 * q < chunk) + { + // Read once: not cached in L1, which then keeps the register spills. + f[k] = __ldcg(reinterpret_cast(rowLogits + 4 * q)); + } + } + // The epoch: every CTA of an exchange launch adds one to the count once it has read it (off the critical path: + // after its final send), and this launch started after the previous one completed, so a CTA reads base + k with + // base a multiple of kCtasPerLaunch and k < kCtasPerLaunch: count / kCtasPerLaunch is the number of earlier + // exchange launches. + unsigned long long* const counter = reinterpret_cast(p.state + kStateCount); + uint64_t const arrivals = *reinterpret_cast(counter); + int32_t const buf = static_cast((arrivals / kCtasPerLaunch) % kNumBuffers); + int32_t const prevBuf = (buf + kNumBuffers - 1) % kNumBuffers; + int32_t const* rec = p.state + kStateRecords + row * kRecordWords; + int32_t const recLevels = rec[0]; + int32_t const recOut = rec[1]; + int32_t const recBnd = rec[2]; + if (t == 0) + { + misc[kMiscRec] = static_cast(recLevels); + misc[kMiscRec + 1] = static_cast(recOut); + misc[kMiscRec + 2] = static_cast(recBnd); + } + if (rowExchange) + { + __syncwarp(); + clusterWait(); + } + + uint32_t validMask = 0; +#pragma unroll + for (int32_t k = 0; k < kVecs; ++k) + { + int32_t const q = t + k * kThreads; + if (rowExchange && 4 * q < chunk) + { + keys[q] = make_uint4(orderedKey(f[k].x), orderedKey(f[k].y), orderedKey(f[k].z), orderedKey(f[k].w)); + int32_t const col = cbase + 4 * q; +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + validMask |= (col + j < v ? 1u : 0u) << (4 * k + j); + } + } + } + stamp(dbg, kDbgLoaded); + + int32_t mode = kModeIdentity; + int32_t levelsUsed = 0; + uint32_t totalSel = 0; + uint32_t nbFinal = 0; + + if (!rowExchange) + { + clearPrevious(p, prevBuf, row, recLevels, recOut, recBnd, wc, c, kCtas); + if (t == 0) + { + atomicAdd(counter, 1ull); + } + clusterSync(); + } + else + { + uint32_t need = kTopK; + uint32_t aboveAllTot = 0; + uint32_t lowTot = 0; + uint32_t prefix = 0; + int32_t level = 0; + for (;; ++level) + { + int32_t const nbins = 1 << levelBits(level); + int32_t const shiftPrev = level == 0 ? 0 : levelShift(level - 1); + int32_t const shiftNow = levelShift(level); + uint32_t const maskNow = (1u << levelBits(level)) - 1u; +#pragma unroll + for (int32_t k = 0; k < kVecs; ++k) + { + uint32_t const vm = (validMask >> (4 * k)) & 0xFu; + if (vm != 0u) + { + uint4 const kv = keys[t + k * kThreads]; + uint32_t const kk[4] = {kv.x, kv.y, kv.z, kv.w}; +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + if (((vm >> j) & 1u) && (level == 0 || (kk[j] >> shiftPrev) == prefix)) + { + atomicAdd(&hist[(kk[j] >> shiftNow) & maskNow], 1u); + } + } + } + } + if (level == 0) + { + stamp(dbg, kDbgLocalHist); + } + __syncthreads(); // this CTA's local histogram of this level is complete + if (level == 0) + { + stamp(dbg, kDbgB0); + } + uint32_t* const ucLevel = p.ucWorkspace + histOff(buf, level, row, 0); + pushAndSend(hist, nbins, slots, histBar0 + 8 * level, ws, p.mcWorkspace + histOff(buf, level, row, rank)); + if (level == 0) + { + stamp(dbg, kDbgSent0); + } + if (t < 2 * kMaxGroups) + { + gsum[t] = 0; // receiveCut's sums, ordered before it by localSuffix's barrier + } + localSuffix(hist, nbins, suff, ws + 192, level == 0 ? misc + kMiscBucketBin : nullptr); + stamp(dbg, kDbgSuffix); + // The next level's local histogram starts from zero (receiveCut's barriers order this before it). + for (int32_t i = t; i < kBins; i += kThreads) + { + hist[i] = 0; + } + if (level == 0) + { + // This CTA's keys in the bins the cut cannot be above, by bin from the top (bin b at + // [suff[b + 1], suff[b])): the keys of a level-0 cut are then prefixes of the buckets. + uint32_t const bucketBin = misc[kMiscBucketBin]; +#pragma unroll + for (int32_t k = 0; k < kVecs; ++k) + { + uint32_t const vm = (validMask >> (4 * k)) & 0xFu; + if (vm != 0u) + { + uint4 const kv = keys[t + k * kThreads]; + uint32_t const kk[4] = {kv.x, kv.y, kv.z, kv.w}; +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + uint32_t const b = kk[j] >> shiftNow; + if (((vm >> j) & 1u) && b >= bucketBin) + { + uint32_t const pos = suff[b + 1] + atomicAdd(&cursor[b], 1u); + bucket[pos] = static_cast(4 * (t + k * kThreads) + j); + } + } + } + } + stamp(dbg, kDbgBucketed); + } + receiveCut(p, ucLevel, suff, nbins, level, need, lowTot, coarse, fine, gsum, misc, pre, preBar0, cut, offs, + wc, 1 + level, dbg); + uint32_t const bin = cut[kCutBin]; + if (bin == kEmpty) + { + abortKernel(wc, 300 + level); + } + putDebug(dbg, kDbgCutGroup, cut[kCutGroup]); + putDebug(dbg, kDbgCutBin, bin); + uint32_t const aboveAll = cut[kCutAboveAll]; + uint32_t const binAll = cut[kCutBinAll]; + uint32_t const aboveLow = cut[kCutAboveLow]; + if (t == 0) + { + misc[kMiscLevelBin + level] = bin; + misc[kMiscLevelGroup + level] = cut[kCutGroup]; + } + aboveAllTot += aboveAll; + lowTot += aboveLow; + need -= aboveAll; + prefix = level == 0 ? bin : ((prefix << levelBits(level)) | bin); + if (need == binAll) + { + mode = kModeFull; + } + else if (level == kLevels - 1) + { + mode = kModeTies; + } + else if (binAll <= static_cast(kBoundaryCap)) + { + mode = kModeCand; + } + else + { + mode = kModeNext; + } + if (mode != kModeNext) + { + break; + } + } + levelsUsed = level + 1; + nbFinal = cut[kCutBinAll]; + putDebug(dbg, kDbgLevels, levelsUsed); + putDebug(dbg, kDbgMode, mode); + putDebug(dbg, kDbgBoundary, nbFinal); + putDebug(dbg, kDbgNeed, need); + + // offs[]: this CTA's offsets, computed by warp 0 in receiveCut at the final cut. + + totalSel = mode == kModeCand ? aboveAllTot : static_cast(kTopK); + int32_t const shiftF = levelShift(level); + uint32_t selCount = 0; + uint32_t bndCount = 0; + uint32_t const outOffC = offs[kOffsOut]; + uint32_t const bndOffC = offs[kOffsBnd]; + uint32_t* const mcOut = p.mcWorkspace + outOff(buf, row) + outOffC; + uint32_t* const mcBnd = p.mcWorkspace + bndOff(buf, row) + 2 * int64_t{bndOffC}; + if (level == 0) + { + // The selected tokens and the boundary pairs are prefixes of the buckets (kModeFull or kModeCand). + uint32_t const nAbove = offs[kOffsAboveCta]; + bndCount = offs[kOffsBndCta]; + selCount = mode == kModeFull ? nAbove + bndCount : nAbove; + stamp(dbg, kDbgStaged); + uint32_t const tokenBase = static_cast(base + cbase); + storeBucketTokens(mcOut, bucket, tokenBase, static_cast(selCount)); + stamp(dbg, kDbgTokSent); + if (mode == kModeCand) + { + storeBucketPairs(mcBnd, bucket + nAbove, reinterpret_cast(keys), tokenBase, + static_cast(bndCount)); + } + } + else + { + // This CTA's keys above the final cut and in its bin. + uint32_t aboveMask = 0; + uint32_t bndMask = 0; +#pragma unroll + for (int32_t k = 0; k < kVecs; ++k) + { + uint32_t const vm = (validMask >> (4 * k)) & 0xFu; + if (vm != 0u) + { + uint4 const kv = keys[t + k * kThreads]; + uint32_t const kk[4] = {kv.x, kv.y, kv.z, kv.w}; +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + uint32_t const s = kk[j] >> shiftF; + if ((vm >> j) & 1u) + { + aboveMask |= (s > prefix ? 1u : 0u) << (4 * k + j); + bndMask |= (s == prefix ? 1u : 0u) << (4 * k + j); + } + } + } + } + stamp(dbg, kDbgClassify); + uint32_t bndBase[kVecs]; + uint32_t selBase[kVecs]; + uint32_t selMask = aboveMask; + if (mode != kModeTies) + { + if (mode == kModeFull) + { + selMask |= bndMask; + } + positionBases(bndMask, selMask, ws, bndBase, selBase, bndCount, selCount); + } + else + { + // The CTA's last offs[kOffsTake] keys of the exact key: the larger tokens. + uint32_t unusedBase[kVecs]; + uint32_t unusedCount = 0; + positionBases(bndMask, 0u, ws, bndBase, unusedBase, bndCount, unusedCount); + uint32_t const firstTaken = offs[kOffsBndCta] - offs[kOffsTake]; +#pragma unroll + for (int32_t k = 0; k < kVecs; ++k) + { + uint32_t pos = bndBase[k]; +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + uint32_t const bit = 1u << (4 * k + j); + if (bndMask & bit) + { + selMask |= pos >= firstTaken ? bit : 0u; + ++pos; + } + } + } + positionBases(selMask, 0u, ws + 128, selBase, unusedBase, selCount, unusedCount); + } + stamp(dbg, kDbgStaged); + + // The selected tokens into every rank's output stage, staged here so that the multicast stores are 16 + // bytes wide; in kModeCand the boundary pairs straight into every rank's boundary stage (8 bytes each). + uint32_t const pairMask = mode == kModeCand ? bndMask : 0u; +#pragma unroll + for (int32_t k = 0; k < kVecs; ++k) + { + if ((((selMask | pairMask) >> (4 * k)) & 0xFu) != 0u) + { + uint4 const kv = keys[t + k * kThreads]; + uint32_t const kk[4] = {kv.x, kv.y, kv.z, kv.w}; + uint32_t sp = selBase[k]; + uint32_t bp = bndBase[k]; +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + uint32_t const bit = 1u << (4 * k + j); + uint32_t const token = static_cast(base + cbase + 4 * (t + k * kThreads) + j); + if (selMask & bit) + { + stageTok[sp] = token; + ++sp; + } + else if (pairMask & bit) + { + st2(mcBnd + 2 * bp, kk[j], token); + } + bp += (bndMask & bit) ? 1u : 0u; + } + } + } + __syncthreads(); + storeWords(mcOut, stageTok, static_cast(selCount)); + } + stamp(dbg, kDbgFinalSent); + if (t == 0) + { + atomicAdd(counter, 1ull); // the epoch (see above), while the stage words travel + } + putDebug(dbg, kDbgSelCount, selCount); + putDebug(dbg, kDbgBndCount, bndCount); + putDebug(dbg, kDbgOutOff, outOffC); + putDebug(dbg, kDbgBndOff, bndOffC); + putDebug(dbg, kDbgOwnAbove, offs[kOffsAboveCta]); + + // The resolver requests all ranks' boundary pairs (nb <= kBoundaryCap: at most two 16-byte words per + // thread) before the output stage is copied. + int32_t const nb = static_cast(nbFinal); + bool const resolver = mode == kModeCand && c == kCtas - 1; + uint32_t const* ucBnd = p.ucWorkspace + bndOff(buf, row); + uint32_t const* ucOut = p.ucWorkspace + outOff(buf, row); + uint4 bv[2] = {make_uint4(0, 0, 0, 0), make_uint4(0, 0, 0, 0)}; + int32_t bvalid[2] = {0, 0}; + if (resolver) + { + int32_t const vecs = (2 * nb + 3) / 4; +#pragma unroll + for (int32_t k = 0; k < 2; ++k) + { + int32_t const vi = t + k * kThreads; + bvalid[k] = vi < vecs ? min(4, 2 * nb - 4 * vi) : 0; + if (bvalid[k] > 0) + { + bv[k] = ldVolatile4(ucBnd + 4 * vi); + } + } + // The level-0 bucket fill counters back to zero: they are the resolve histogram. + for (int32_t i = static_cast(misc[kMiscBucketBin]) + t; i < kBins; i += kThreads) + { + cursor[i] = 0; + } + __syncthreads(); + } + else + { + // What the previous exchange launch wrote into buffer prevBuf for this row, cleared by the CTAs other + // than the resolver while the peers' stage words travel. + clearPrevious(p, prevBuf, row, static_cast(misc[kMiscRec]), + static_cast(misc[kMiscRec + 1]), static_cast(misc[kMiscRec + 2]), wc, c, + mode == kModeCand ? kCtas - 1 : kCtas); + stamp(dbg, kDbgCleared); + } + // The output stage of every rank into the output row: this CTA's kOutPerCta positions. + if (t < kOutPerCta / 4) + { + int32_t const pos = c * kOutPerCta + 4 * t; + int32_t const valid = min(4, max(0, static_cast(totalSel) - pos)); + if (valid > 0) + { + uint4 const w = pollVec(ucOut + pos, valid, wc, 10); + if (valid == 4) + { + *reinterpret_cast(outRow + pos) = make_int4(static_cast(w.x), + static_cast(w.y), static_cast(w.z), static_cast(w.w)); + } + else + { + uint32_t const ww[4] = {w.x, w.y, w.z, w.w}; + for (int32_t j = 0; j < valid; ++j) + { + outRow[pos + j] = static_cast(ww[j]); + } + } + } + } + stamp(dbg, kDbgOutCopied); + if (resolver) + { + PollClock clk; + while (anyEmpty(bv[0], bvalid[0]) || anyEmpty(bv[1], bvalid[1])) + { + clk.tick(wc, 11); +#pragma unroll + for (int32_t k = 0; k < 2; ++k) + { + if (anyEmpty(bv[k], bvalid[k])) + { + bv[k] = ldVolatile4(ucBnd + 4 * (t + k * kThreads)); + } + } + } + stamp(dbg, kDbgBndPolled); + uint32_t const pk[4] = {bv[0].x, bv[0].z, bv[1].x, bv[1].z}; + uint32_t const pt[4] = {bv[0].y, bv[0].w, bv[1].y, bv[1].w}; + uint32_t const pv = (bvalid[0] >= 2 ? 1u : 0u) | (bvalid[0] == 4 ? 2u : 0u) | (bvalid[1] >= 2 ? 4u : 0u) + | (bvalid[1] == 4 ? 8u : 0u); + resolveBoundary(pk, pt, pv, need, shiftF, prefix << shiftF, rhist, ws + 192, cut, small, + misc + kMiscSmallCount, outRow, static_cast(totalSel), dbg); + stamp(dbg, kDbgResolved); + } + } + + // This launch's record for the next exchange launch's clearing (every CTA of the row read the previous record + // before it pushed its level-0 counts, which CTA 0 waited for; a row without an exchange passed its cluster + // barrier). + if (t == 0 && c == 0) + { + int32_t* recW = p.state + kStateRecords + row * kRecordWords; + recW[0] = levelsUsed; + recW[1] = static_cast(totalSel); + recW[2] = mode == kModeCand ? static_cast(nbFinal) : 0; + } + stamp(dbg, kDbgEnd); + putDebug(dbg, kDbgClockEnd, static_cast(clock64())); + flushDebug(p.debug, row, c); +#else + (void) p; +#endif +} + +__global__ void dsaHistShardMetaKernel(int32_t const* kvLens, int32_t const* blockTable, int32_t blockTableWidth, + int32_t worldSize, int32_t rank, int32_t tokensPerBlock, int32_t* shardCtx, int32_t* shardBlockTable, + int32_t shardBlockTableWidth) +{ + int32_t const kvLen = kvLens[0]; + int32_t const perRank = (max(kvLen, 0) + worldSize - 1) / worldSize; + int32_t const size = (perRank + kShardAlign - 1) / kShardAlign * kShardAlign; + int32_t const base = rank * size; + int32_t const n = min(max(kvLen - base, 0), size); + if (threadIdx.x == 0) + { + shardCtx[0] = max(n, 1); + } + int32_t const firstPage = base / tokensPerBlock; + int32_t const pages = (n + tokensPerBlock - 1) / tokensPerBlock; + for (int32_t j = threadIdx.x; j < shardBlockTableWidth; j += blockDim.x) + { + shardBlockTable[j] = (j < pages && firstPage + j < blockTableWidth) ? blockTable[firstPage + j] : 0; + } +} + +int64_t defaultTimeoutNs() +{ + return kDefaultTimeoutMs * 1000000; +} + +} // namespace + +void invokeDsaHistTopK(HistTopKParams const& inParams, cudaStream_t stream) +{ + HistTopKParams params = inParams; + if (params.numRows <= 0 || params.numRows > kMaxRows || params.worldSize <= 0 || params.worldSize > kMaxRanks + || params.rank < 0 || params.rank >= params.worldSize || params.localCapacity > kMaxShard + || params.logitsStride % 4 != 0 || reinterpret_cast(params.logits) % 16 != 0 + || reinterpret_cast(params.state) % 8 != 0) + { + throw std::invalid_argument("glm_dsa_hist_topk: unsupported rows, group, shard capacity or alignment"); + } + if (params.timeoutNs == 0) + { + params.timeoutNs = static_cast(defaultTimeoutNs()); + } + static thread_local bool configured[64] = {}; + int device = 0; + cudaGetDevice(&device); + cudaError_t err = cudaSuccess; + if (device >= 64 || !configured[device]) + { + err = cudaFuncSetAttribute( + dsaHistTopKKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(kSmemBytes)); + if (err != cudaSuccess) + { + throw std::runtime_error(std::string("glm_dsa_hist_topk: smem attribute: ") + cudaGetErrorString(err)); + } + if (device < 64) + { + configured[device] = true; + } + } + cudaLaunchConfig_t config{}; + config.gridDim = dim3(kCtas, kMaxRows); + config.blockDim = dim3(kThreads); + config.dynamicSmemBytes = kSmemBytes; + config.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kCtas; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL() ? 1 : 0; + config.attrs = attrs; + config.numAttrs = 2; + err = cudaLaunchKernelEx(&config, dsaHistTopKKernel, params); + if (err == cudaSuccess) + { + err = cudaGetLastError(); + } + if (err != cudaSuccess) + { + throw std::runtime_error(std::string("glm_dsa_hist_topk: launch: ") + cudaGetErrorString(err)); + } +} + +void invokeDsaHistShardMeta(int32_t const* kvLens, int32_t const* blockTable, int32_t blockTableWidth, + int32_t worldSize, int32_t rank, int32_t tokensPerBlock, int32_t* shardCtx, int32_t* shardBlockTable, + int32_t shardBlockTableWidth, cudaStream_t stream) +{ + if (worldSize <= 0 || worldSize > kMaxRanks || rank < 0 || rank >= worldSize || tokensPerBlock <= 0 + || kShardAlign % tokensPerBlock != 0) + { + throw std::invalid_argument("glm_dsa_hist_shard_meta: unsupported group or page size"); + } + dsaHistShardMetaKernel<<<1, 256, 0, stream>>>(kvLens, blockTable, blockTableWidth, worldSize, rank, + tokensPerBlock, shardCtx, shardBlockTable, shardBlockTableWidth); + cudaError_t const err = cudaGetLastError(); + if (err != cudaSuccess) + { + throw std::runtime_error(std::string("glm_dsa_hist_shard_meta: launch: ") + cudaGetErrorString(err)); + } +} + +} // namespace kernels::glm_moe_dsa::hist_topk + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.h new file mode 100644 index 000000000000..a50ddc363457 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.h @@ -0,0 +1,142 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::hist_topk +{ + +// Key-sharded DSA indexer top-k of one generation request, exact, by histogram exchange. +// +// The indexer context of the request is split over the ranks of a tensor-parallel group into contiguous shards: +// S = ceil(ceil(kvLen / worldSize) / kShardAlign) * kShardAlign, rank r holds the tokens [r S, r S + n_r) with +// n_r = clamp(kvLen - r S, 0, S), and scores only those (the caller supplies the shard's logits). Verify row i of +// nextN sees the tokens [0, kvLen - nextN + i + 1). +// +// Per row, one cluster of kCtasPerRow CTAs holds the row's shard keys (order-preserving keys of the FP32 logits). +// Every rank +// 1. builds the histogram of its keys' top 11 bits (per CTA in shared memory, reduced over the cluster: every CTA +// owns 1/kCtasPerRow of the bins and receives the other CTAs' counts of them as asynchronous stores into its +// shared memory, which complete on its barrier) and stores it once through the multicast address: the 32-bit +// totals of its 32-bin groups, and two 16-bit bin counts per word, saturated at 0xFFFF (at most one bin of a +// group can saturate, as two would hold more than kMaxShard keys; the receivers recover its count from the group +// total); +// 2. reads all ranks' group totals, finds the group that holds the kTopK-th largest key, reads that group's bins +// of all ranks and finds the cut bin: the keys above it are selected, and every rank knows each rank's count +// above it and in it; +// 3. when that bin is not taken whole and holds more than kBoundaryCap keys, refines the same way with the next +// 11 bits of the keys in the bin, then with the last 10 bits (the exact key, whose ties are taken from the +// higher ranks first: the larger tokens first); +// 4. stores its selected tokens through the multicast address into every rank's output stage, at its offset (rank, +// then CTA), and, when the boundary bin is neither taken whole nor the exact key, its boundary (key, token) pairs +// into every rank's boundary stage. After a cut on the top 11 bits (the common case) a CTA's tokens and pairs are +// prefixes of buckets it fills while it waits for the peers' histograms: its keys in the bins that cut cannot be +// above, by bin from the top, in the order of shared-memory atomics within a bin; after a later cut they are in +// position order; +// 5. copies the output stage into the output row; one CTA selects the boundary pairs that complete the row, by key +// and then the larger token, and appends them in stage order. +// Every rank derives each decision from the same multicast data, so all ranks write the same bits: the exact top-kTopK +// of the row by (logit, token), the larger token first among equal logits. The order within the row is the stages' +// order: the same on all ranks, but after a cut on the top 11 bits it can differ between calls with the same logits. +// Rows of at most kTopK tokens get [0, n) and -1 without an exchange, and a launch whose rows all have at most kTopK +// tokens touches neither the workspace nor the state. +// +// Workspace: kNumBuffers Lamport buffers, rotated by a device epoch that counts exchange launches. In exchange launch x +// a rank clears, on its own copy, what launch x - 1 wrote into buffer (x - 1) % kNumBuffers (the sizes are recorded +// in the state per row), each histogram word after it reads the word written (launch x - 1 read only the group +// totals and one group's bins); a peer writes that buffer again only in launch x + 2, after this rank completed +// launch x + 1. +// +// A wait for a peer's word that exceeds the timeout records the wait in the state's error word and aborts the kernel: +// the selection can no longer complete, and a late peer write could land in a buffer this rank reuses. + +constexpr int32_t kTopK = 2048; +constexpr int32_t kMaxRanks = 16; +constexpr int32_t kMaxRows = 8; +//! CTAs per row (one cluster). +constexpr int32_t kCtasPerRow = 8; +constexpr int32_t kShardAlign = 128; +//! Largest shard (tokens per rank): three float4 per thread of kCtasPerRow CTAs of 1024 threads. At 16 ranks a request +//! of up to 1,572,864 indexer tokens. +constexpr int32_t kMaxShard = 98304; +//! Boundary-bin keys exchanged per row; a fuller boundary bin is refined with the next digit instead. +constexpr int32_t kBoundaryCap = 4096; +constexpr int32_t kNumBuffers = 3; +//! Bins of the first two digits (11 bits); the last digit has 10 bits. +constexpr int32_t kBins = 2048; +//! Bins per group of the coarse summary. +constexpr int32_t kGroupBins = 32; +//! Lamport words of one (level, row, rank) histogram: kBins / 2 words of two 16-bit bin counts, then the 32-bit totals +//! of the kBins / kGroupBins groups. +constexpr int32_t kHistWords = kBins / 2 + kBins / kGroupBins; +constexpr int32_t kLevels = 3; +constexpr int64_t kBufferWords = int64_t{kLevels} * kMaxRows * kMaxRanks * kHistWords + int64_t{kMaxRows} * kTopK + + int64_t{kMaxRows} * kBoundaryCap * 2; +//! 32-bit words of the exchange workspace; every word must start as 0xFFFFFFFF. +constexpr int64_t kWorkspaceWords = kNumBuffers * kBufferWords; +//! int32 words of the per-rank device state (zero-initialized): a 64-bit count of the CTAs that completed exchange +//! launches (the epoch is that count / (kCtasPerRow * kMaxRows)), the error word, a spare word, then per row the +//! levels, output-stage words and boundary pairs the last exchange launch wrote. +constexpr int32_t kStateWords = 4 + kMaxRows * 4; +//! Wait for a peer's words before the kernel aborts, when HistTopKParams::timeoutNs is 0. +constexpr int64_t kDefaultTimeoutMs = 120000; +//! Debug slots per CTA (phase %globaltimer stamps and per-row statistics), for diagnosis only. +constexpr int32_t kDebugSlots = 64; + +struct HistTopKParams +{ + //! [numRows, logitsStride] FP32 logits of this rank's shard (columns [0, localCapacity), localCapacity >= S). + float const* logits; + int64_t logitsStride; + int32_t localCapacity; + //! [1] int32 indexer KV length of the request. + int32_t const* kvLens; + //! [numRows, kTopK] int32 output: request-local token indices, -1 for unused slots. + int32_t* out; + int32_t numRows; + int32_t nextN; + int32_t rank; + int32_t worldSize; + //! Multicast and this rank's unicast address of the exchange workspace (kWorkspaceWords words). + uint32_t* mcWorkspace; + uint32_t* ucWorkspace; + //! [kStateWords] int32 per-rank device state, 8-byte aligned. + int32_t* state; + //! Optional [kMaxRows, kCtasPerRow, kDebugSlots] int64 (nullptr = off). + int64_t* debug = nullptr; + //! Longest wait for a peer's words (ns); 0 = the default. + uint64_t timeoutNs = 0; +}; + +void invokeDsaHistTopK(HistTopKParams const& params, cudaStream_t stream); + +//! Shard of the generation request for `rank`: shardCtx[0] = max(n_r, 1) (an empty shard scores one token no row +//! reaches) and shardBlockTable[j] = blockTable[r S / tokensPerBlock + j] for the shard's pages, 0 beyond. One CTA. +void invokeDsaHistShardMeta(int32_t const* kvLens, int32_t const* blockTable, int32_t blockTableWidth, + int32_t worldSize, int32_t rank, int32_t tokensPerBlock, int32_t* shardCtx, int32_t* shardBlockTable, + int32_t shardBlockTableWidth, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa::hist_topk + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index e9971fb08213..b649ddc4a283 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -134,6 +134,7 @@ add_th_common_family( allgatherOp.cpp allreduceOp.cpp alltoallOp.cpp + glmDsaHistTopKOp.cpp asyncUlyssesOp.cpp glmMnnvlAllgatherOp.cpp glmMnnvlHopOp.cpp diff --git a/cpp/tensorrt_llm/thop/glmDsaHistTopKOp.cpp b/cpp/tensorrt_llm/thop/glmDsaHistTopKOp.cpp new file mode 100644 index 000000000000..fb0ab897d100 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmDsaHistTopKOp.cpp @@ -0,0 +1,120 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/mcastDevMemUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.h" +#include "tensorrt_llm/runtime/mcastDeviceMemory.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include + +#include + +namespace tkh = tensorrt_llm::kernels::glm_moe_dsa::hist_topk; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +// Exact top-2048 of one generation request's key-sharded GLM DSA indexer logits (glmDsaHistTopK.h): this rank's shard +// logits [num_rows, >= shard size] FP32 in, the request-local token indices [num_rows, 2048] int32 out, identical on +// every rank of the group. comm_buffer is this rank's unicast view of a McastGPUBuffer workspace of +// kWorkspaceWords int32 words, all -1 at creation; state holds kStateWords int32, zero at creation; both are private +// to this op. +void glmDsaHistTopK(at::Tensor const& logits, at::Tensor const& kvLens, at::Tensor& out, at::Tensor& commBuffer, + at::Tensor& state, int64_t nextN, int64_t groupSize) +{ + auto* mcastMem = tensorrt_llm::common::findMcastDevMemBuffer(commBuffer.data_ptr()); + TORCH_CHECK(mcastMem != nullptr, "glm_dsa_hist_topk: comm_buffer must come from a McastGPUBuffer"); + TORCH_CHECK(static_cast(mcastMem->getWorldSize()) == groupSize, "glm_dsa_hist_topk: group_size ", + groupSize, " does not match the workspace group of ", mcastMem->getWorldSize(), " ranks"); + CHECK_TH_CUDA(logits); + CHECK_TH_CUDA(kvLens); + CHECK_TH_CUDA(out); + CHECK_TH_CUDA(state); + TORCH_CHECK(logits.scalar_type() == at::ScalarType::Float && logits.dim() == 2 && logits.stride(1) == 1, + "glm_dsa_hist_topk: logits must be a row-major FP32 matrix"); + TORCH_CHECK(kvLens.scalar_type() == at::ScalarType::Int && kvLens.numel() >= 1, "glm_dsa_hist_topk: kv_lens int32"); + TORCH_CHECK(out.scalar_type() == at::ScalarType::Int && out.dim() == 2 && out.size(1) == tkh::kTopK + && out.stride(1) == 1 && out.stride(0) == tkh::kTopK, + "glm_dsa_hist_topk: out must be a contiguous int32 [num_rows, 2048] matrix"); + int64_t const numRows = out.size(0); + TORCH_CHECK(numRows >= 1 && numRows <= tkh::kMaxRows && logits.size(0) >= numRows && nextN >= numRows, + "glm_dsa_hist_topk: 1..8 rows of one request"); + TORCH_CHECK(logits.size(1) <= tkh::kMaxShard, "glm_dsa_hist_topk: at most ", tkh::kMaxShard, " shard columns"); + TORCH_CHECK(commBuffer.scalar_type() == at::ScalarType::Int && commBuffer.numel() >= tkh::kWorkspaceWords, + "glm_dsa_hist_topk: comm_buffer must hold ", tkh::kWorkspaceWords, " int32 words"); + TORCH_CHECK(state.scalar_type() == at::ScalarType::Int && state.numel() >= tkh::kStateWords, + "glm_dsa_hist_topk: state must hold ", tkh::kStateWords, " int32 words"); + + tkh::HistTopKParams params; + params.logits = logits.data_ptr(); + params.logitsStride = logits.stride(0); + params.localCapacity = static_cast(logits.size(1)); + params.kvLens = kvLens.data_ptr(); + params.out = out.data_ptr(); + params.numRows = static_cast(numRows); + params.nextN = static_cast(nextN); + params.rank = static_cast(mcastMem->getRank()); + params.worldSize = static_cast(groupSize); + params.mcWorkspace = static_cast(mcastMem->getMulticastPtr()); + params.ucWorkspace = reinterpret_cast(commBuffer.data_ptr()); + params.state = state.data_ptr(); + tkh::invokeDsaHistTopK(params, at::cuda::getCurrentCUDAStream(logits.get_device())); +} + +// Shard metadata of the generation request for this rank (glmDsaHistTopK.h): shard_ctx [1] and shard_block_table +// [1, width] int32, from the request's KV length and block table row. +void glmDsaHistShardMeta(at::Tensor const& kvLens, at::Tensor const& blockTable, at::Tensor& shardCtx, + at::Tensor& shardBlockTable, int64_t groupSize, int64_t rank, int64_t tokensPerBlock) +{ + CHECK_TH_CUDA(kvLens); + CHECK_TH_CUDA(blockTable); + TORCH_CHECK(kvLens.scalar_type() == at::ScalarType::Int && kvLens.numel() >= 1, + "glm_dsa_hist_shard_meta: kv_lens must be int32"); + TORCH_CHECK(blockTable.scalar_type() == at::ScalarType::Int && blockTable.dim() >= 1 && blockTable.stride(-1) == 1, + "glm_dsa_hist_shard_meta: block_table must be int32 with contiguous rows"); + TORCH_CHECK(shardCtx.scalar_type() == at::ScalarType::Int && shardCtx.numel() >= 1, + "glm_dsa_hist_shard_meta: shard_ctx must be int32"); + TORCH_CHECK(shardBlockTable.scalar_type() == at::ScalarType::Int && shardBlockTable.is_contiguous(), + "glm_dsa_hist_shard_meta: shard_block_table must be contiguous int32"); + tkh::invokeDsaHistShardMeta(kvLens.data_ptr(), blockTable.data_ptr(), + static_cast(blockTable.size(-1)), static_cast(groupSize), static_cast(rank), + static_cast(tokensPerBlock), shardCtx.data_ptr(), shardBlockTable.data_ptr(), + static_cast(shardBlockTable.numel()), at::cuda::getCurrentCUDAStream(kvLens.get_device())); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_dsa_hist_topk(Tensor logits, Tensor kv_lens, Tensor(a!) out, Tensor(b!) comm_buffer, Tensor(c!) state, " + "int next_n, int group_size) -> ()"); + m.def( + "glm_dsa_hist_shard_meta(Tensor kv_lens, Tensor block_table, Tensor(a!) shard_ctx, " + "Tensor(b!) shard_block_table, int group_size, int rank, int tokens_per_block) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_dsa_hist_topk", &tensorrt_llm::torch_ext::glmDsaHistTopK); + m.impl("glm_dsa_hist_shard_meta", &tensorrt_llm::torch_ext::glmDsaHistShardMeta); +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_hist_topk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_hist_topk.py new file mode 100644 index 000000000000..a9973f360be2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_hist_topk.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shard geometry and exchange workspace of the key-sharded decode indexer top-k. + +The decode indexer is replicated over the tensor-parallel group: without sharding every rank +computes the same FP8 MQA logits over the whole context and the same top-2048. Instead rank ``r`` +scores only its contiguous shard of the request's indexer keys (``glm_dsa_hist_shard_meta`` builds +the shard's context length and block table), and ``glm_dsa_hist_topk`` finds every row's exact +top-2048 over all shards by exchanging key histograms through a private multicast workspace; every +rank ends with the same rows. See ``cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.h``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Dict, Optional + +import torch + +from tensorrt_llm.mapping import Mapping + +# Constants of glmDsaHistTopK.h. +MAX_RANKS = 16 +SHARD_ALIGN = 128 +MAX_SHARD = 98304 +MAX_ROWS = 8 +TOPK = 2048 +WORKSPACE_WORDS = 3 * (3 * 8 * 16 * 1088 + 8 * 2048 + 8 * 4096 * 2) +STATE_WORDS = 4 + 8 * 4 + + +def shard_size(kv_len: int, group_size: int) -> int: + """Keys per rank of a request of ``kv_len`` indexer keys: ceil(kv_len / group_size), rounded + up to SHARD_ALIGN. Rank ``r`` holds the keys [r * S, r * S + n_r), n_r = clamp(kv_len - r * S, + 0, S).""" + per_rank = -(-max(kv_len, 0) // group_size) + return -(-per_rank // SHARD_ALIGN) * SHARD_ALIGN + + +def shard_capacity(max_seq_len: int, group_size: int) -> int: + """Largest shard of a request of at most ``max_seq_len`` indexer keys.""" + return shard_size(max_seq_len, group_size) + + +@dataclass +class HistTopkWorkspace: + """The exchange workspace and device state of ``glm_dsa_hist_topk`` for one TP group. + + ``comm_buffer`` is this rank's unicast view of ``WORKSPACE_WORDS`` int32 words of a + McastGPUBuffer, all -1 at creation; ``state`` is ``STATE_WORDS`` int32 words, zero at creation. + Both are private to the op and live as long as the engine: captured CUDA graphs launch on + them. ``handle`` keeps the multicast allocation alive. + """ + + group_size: int + handle: Any + comm_buffer: torch.Tensor + state: torch.Tensor + comm: Any + + +def hist_topk_supported(mapping: Mapping) -> bool: + """Whether the key-sharded top-k can serve this mapping: a pure tensor-parallel group of 2 to + 16 ranks without attention DP or context parallelism, on an MNNVL fabric.""" + from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce + + return ( + 2 <= mapping.tp_size <= MAX_RANKS + and mapping.world_size == mapping.tp_size + and not mapping.has_cp() + and not mapping.enable_attention_dp + and MNNVLAllReduce.is_mnnvl(mapping, torch.bfloat16) + ) + + +def create_hist_topk_workspace(mapping: Mapping, comm: Optional[Any] = None) -> HistTopkWorkspace: + """Allocate the workspace on every rank of ``mapping``'s TP group (a collective call). + + ``comm`` is the TP-group communicator the multicast handles are exchanged over: the MNNVL + all-reduce's (an MPI split, or the TP ProcessGroup without MPI) when None. The outcome is + converged over the group, so every rank raises or none does. + """ + from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, + _make_mnnvl_mcast_buffer, + _mnnvl_device_index, + _mnnvl_workspace_all_succeeded, + ) + + if comm is None: + comm = _get_mnnvl_workspace_comm(mapping) + error: Optional[Exception] = None + workspace = None + try: + handle = _make_mnnvl_mcast_buffer( + comm, WORKSPACE_WORDS * 4, mapping, use_fabric_handle=mapping.is_multi_node() + ) + comm_buffer = handle.get_uc_buffer(mapping.tp_rank, (WORKSPACE_WORDS,), torch.int32, 0) + with torch.inference_mode(): + comm_buffer.fill_(-1) + state = torch.zeros( + STATE_WORDS, + dtype=torch.int32, + device=torch.device("cuda", _mnnvl_device_index(mapping)), + ) + torch.cuda.synchronize() + workspace = HistTopkWorkspace(mapping.tp_size, handle, comm_buffer, state, comm) + except Exception as e: + # Any failure is converged over the group and re-raised on every rank below: a rank that + # raised alone would leave its peers waiting in the next exchange. + error = e + if not _mnnvl_workspace_all_succeeded(comm, error is None): + raise RuntimeError( + "GLM DSA histogram top-k workspace construction failed on at least one rank" + ) from error + return workspace + + +_WORKSPACES: Dict[Mapping, HistTopkWorkspace] = {} + + +def get_hist_topk_workspace(mapping: Mapping) -> HistTopkWorkspace: + """The workspace of ``mapping``'s TP group, created on the first call (collective: every rank + of the group makes its first call in the same place, outside CUDA graph capture).""" + workspace = _WORKSPACES.get(mapping) + if workspace is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("the GLM DSA histogram top-k workspace is created outside capture") + workspace = _WORKSPACES[mapping] = create_hist_topk_workspace(mapping) + return workspace diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/hist_shard_meta.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/hist_shard_meta.py new file mode 100644 index 000000000000..f497903c8adf --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/hist_shard_meta.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""This rank's key shard of the generation request: its context length and block table.""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._hist_topk import shard_capacity, shard_size + +#: c1: 16 ranks, 64-token pages, and the shard table of the gen server's indexer capacity. +_C1 = dict(group_size=16, tokens_per_block=64, capacity=shard_capacity(1_048_596, 16)) + + +def _spec(kv_len, rank, **overrides) -> dict: + spec = dict(_C1, kv_len=kv_len, rank=rank) + spec.update(overrides) + return spec + + +class GlmHistShardMeta(OpWrapper): + """``shard_ctx[0]`` and ``shard_block_table[0, :]`` of rank ``rank`` for one request. + + With ``S = shard_size(kv_lens[0], group_size)`` (ceil(kv_len / group_size) rounded up to 128), + the rank holds the keys [rank * S, rank * S + n), ``n = clamp(kv_len - rank * S, 0, S)``: + + - ``shard_ctx[0] = max(n, 1)`` (an empty shard scores one key no row reaches); + - ``shard_block_table[0, j] = block_table[rank * S / tokens_per_block + j]`` for the shard's + ``ceil(n / tokens_per_block)`` pages, 0 for the rest of the row. + + ``block_table`` is the request's INDEX_KEY block-table row (one-dimensional). One CTA. Returns + ``(shard_ctx, shard_block_table)``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="c1: a 530,000-key request on rank 0", spec=_spec(530_000, 0)), + Cell(why="c1: the same request's last rank, a partial shard", spec=_spec(530_000, 15)), + Cell(why="the full 1,048,596-key capacity on rank 15", spec=_spec(1_048_596, 15)), + Cell(why="a short request: rank 13 holds no key, one unread key", spec=_spec(3_000, 13)), + Cell( + why="a shard ending inside a page: 1,000,003 keys on rank 15", + spec=_spec(1_000_003, 15), + ), + ) + + note = """ + Metadata consumed: none. The target passes the generation request's indexer KV length (int32, + on the device), its row of the INDEX_KEY block offsets, persistent ``shard_ctx`` [1] and + ``shard_block_table`` [1, ceil(capacity / 64)] int32 buffers (captured CUDA graphs read + them; capacity = ``_hist_topk.shard_capacity(indexer max_seq_len, 16)``), the group size + (16), this rank's TP rank and the page size (64). + + The kernel refuses a group above 16 ranks and a page size that does not divide 128. + """ + + def __call__( + self, + kv_lens: torch.Tensor, + block_table: torch.Tensor, + shard_ctx: torch.Tensor, + shard_block_table: torch.Tensor, + group_size: int, + rank: int, + tokens_per_block: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + torch.ops.trtllm.glm_dsa_hist_shard_meta( + kv_lens, block_table, shard_ctx, shard_block_table, group_size, rank, tokens_per_block + ) + return shard_ctx, shard_block_table + + def reference( + self, + kv_lens: torch.Tensor, + block_table: torch.Tensor, + shard_ctx: torch.Tensor, + shard_block_table: torch.Tensor, + group_size: int, + rank: int, + tokens_per_block: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Fresh ``(shard_ctx, shard_block_table)`` computed on the host.""" + kv_len = int(kv_lens[0]) + size = shard_size(kv_len, group_size) + base = rank * size + n = min(max(kv_len - base, 0), size) + first_page = base // tokens_per_block + pages = -(-n // tokens_per_block) + row = block_table.tolist() + width = shard_block_table.numel() + table = [ + row[first_page + j] if j < pages and first_page + j < len(row) else 0 + for j in range(width) + ] + ctx = torch.tensor([max(n, 1)], dtype=torch.int32, device=shard_ctx.device) + out = torch.tensor(table, dtype=torch.int32, device=shard_block_table.device) + return ctx, out.view(shard_block_table.shape) + + def is_valid( + self, + kv_lens: torch.Tensor, + block_table: torch.Tensor, + shard_ctx: torch.Tensor, + shard_block_table: torch.Tensor, + group_size: int, + rank: int, + tokens_per_block: int, + ) -> None: + # The kernel reads the first row of a two-dimensional table: another request's pages. + assert block_table.dim() == 1, "block_table must be the request's one-dimensional row" + # A table narrower than the shard's pages drops pages the logits then read as page 0. + kv_len = int(kv_lens[0]) + size = shard_size(kv_len, group_size) + n = min(max(kv_len - rank * size, 0), size) + assert shard_block_table.numel() * tokens_per_block >= n, ( + f"shard_block_table holds {shard_block_table.numel()} pages, the shard {n} keys" + ) + # Pages past the request's row are read as 0 instead of failing. + if n > 0: + first_page = rank * size // tokens_per_block + assert first_page + -(-n // tokens_per_block) <= block_table.numel(), ( + "the shard's pages run past the request's block-table row" + ) + + +hist_shard_meta = GlmHistShardMeta() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/hist_topk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/hist_topk.py new file mode 100644 index 000000000000..3051f1d2f81f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/hist_topk.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Exact top-2048 of the key-sharded decode indexer logits, by an MNNVL histogram exchange.""" + +from __future__ import annotations + +from typing import Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._hist_topk import MAX_SHARD, TOPK, shard_size + +_TOKEN_BITS = 21 + + +def _spec(group_size, kv_len, pattern="randn", rows=8, next_n=8) -> dict: + return dict(group_size=group_size, kv_len=kv_len, pattern=pattern, rows=rows, next_n=next_n) + + +def canonical_rows(rows: torch.Tensor) -> torch.Tensor: + """Each row's tokens in ascending order, the -1 slots last: the set a row holds, as a tensor.""" + largest = torch.iinfo(rows.dtype).max + keyed = torch.where(rows < 0, torch.full_like(rows, largest), rows) + ordered = keyed.sort(dim=1).values + return torch.where(ordered == largest, torch.full_like(ordered, -1), ordered) + + +def ordered_keys(logits: torch.Tensor) -> torch.Tensor: + """The kernel's order-preserving 32-bit key of each FP32 logit, as int64: the bits with the + sign bit set when it is clear and all bits inverted when it is set, so -0 sorts below +0, + positive NaNs above +inf and negative NaNs below -inf; the one key equal to the empty word + (NaN 0x7FFFFFFF) is clamped one below it.""" + bits = logits.contiguous().view(torch.int32).long() & 0xFFFFFFFF + keys = torch.where(bits >= 0x80000000, (~bits) & 0xFFFFFFFF, bits | 0x80000000) + return keys.clamp(max=0xFFFFFFFE) + + +class GlmHistTopk(OpWrapper): + """Each verify row's exact top-2048 over the TP group's key shards; every rank gets the same + rows. + + The request of ``kv_lens[0]`` indexer keys is split into contiguous shards of ``S = + shard_size(kv_len, group_size)`` keys, rank ``r`` holding [r * S, r * S + n_r); ``logits`` + [num_rows, >= S] fp32 holds this rank's shard (its column ``c`` is key r * S + c). Verify row + ``i`` of ``next_n`` sees the keys [0, kv_len - next_n + i + 1). + + ``out[i]`` [2048] int32 gets the row's top-2048 keys by (logit, key), the larger key first among + equal logits, in an unspecified order that is the same on every rank; a row of at most 2048 + keys gets [0, n) and then -1. FP32 values are ordered by their bits: -0 below +0, NaNs above + +inf (positive) or below -inf (negative). + + The exchange runs through ``comm_buffer`` (this rank's unicast view of the group's multicast + workspace) and ``state`` (``_hist_topk.HistTopkWorkspace``), both private to the op: every rank + of the group makes the same sequence of calls. Returns ``out``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="16 ranks (c1): a 530,000-key request, 8 verify rows", spec=_spec(16, 530_000)), + Cell( + why="16 ranks: the c1 gen server's 1,048,596-key capacity", + spec=_spec(16, 1_048_596), + ), + Cell( + why="16 ranks: the largest shard, 98,304 keys of a 1,572,864-key request", + spec=_spec(16, 16 * MAX_SHARD), + ), + Cell(why="4 ranks: a 300,000-key request, random logits", spec=_spec(4, 300_000)), + Cell(why="4 ranks: ties at the cut, the larger keys taken", spec=_spec(4, 200_000, "ties")), + Cell(why="4 ranks: every logit equal", spec=_spec(4, 50_000, "equal")), + Cell(why="4 ranks: +-inf, +-0 and NaN among the logits", spec=_spec(4, 120_000, "special")), + Cell( + why="4 ranks: one histogram bin holds more than 65,535 keys of a shard", + spec=_spec(4, 4 * MAX_SHARD, "saturated"), + ), + Cell( + why="4 ranks: 2,051 keys, rows 0-4 without an exchange, rows 5-7 with", + spec=_spec(4, 2_051), + ), + Cell(why="4 ranks: no row exceeds 2,048 keys", spec=_spec(4, 1_000)), + Cell(why="4 ranks: one row of a next_n = 1 call", spec=_spec(4, 90_000, rows=1, next_n=1)), + ) + + note = """ + Metadata consumed: none. At c1 the target passes the shard logits of + ``fp8_paged_mqa_logits`` (the view ``logits[:, :shard capacity]``), the generation request's + indexer KV length ([1] int32 on the device), its [8, 2048] TopK buffer, the workspace's + ``comm_buffer`` and ``state``, next_n = 8 and the TP size (16). The workspace is created once + per process (``_hist_topk.get_hist_topk_workspace``) by every rank of the group, outside CUDA + graph capture. + + Order and ties: the set is exact; the order within a row after a cut on the top 11 key bits + follows shared-memory atomics and can change between calls with the same logits. + + The launch is PDL-enabled (``TRTLLM_ENABLE_PDL``): before its wait it initializes its cluster + barriers and reads ``kv_lens`` (written by the step's metadata preparation, before the logits). + A wait for a peer longer than 120 s records the error in ``state`` and aborts the kernel. + + The kernel refuses (loudly) a shard above 98,304 keys, more than 8 rows or 16 ranks, logits rows + not 16-byte aligned, and a workspace of another group. + """ + + def __call__( + self, + logits: torch.Tensor, + kv_lens: torch.Tensor, + out: torch.Tensor, + comm_buffer: torch.Tensor, + state: torch.Tensor, + next_n: int, + group_size: int, + ) -> torch.Tensor: + torch.ops.trtllm.glm_dsa_hist_topk( + logits, kv_lens, out, comm_buffer, state, next_n, group_size + ) + return out + + def reference( + self, + logits: Sequence[torch.Tensor], + kv_lens: torch.Tensor, + out: torch.Tensor, + comm_buffer: torch.Tensor, + state: torch.Tensor, + next_n: int, + group_size: int, + ) -> torch.Tensor: + """``canonical_rows`` of the expected ``out``, from the group's shard logits in rank order. + The workspace only carries the exchange.""" + assert len(logits) == group_size, "the reference takes every rank's shard logits" + kv_len = int(kv_lens[0]) + size = shard_size(kv_len, group_size) + rows = out.shape[0] + expected = torch.full((rows, TOPK), -1, dtype=torch.int32, device=out.device) + whole = [] + for rank, shard in enumerate(logits): + n = min(max(kv_len - rank * size, 0), size) + whole.append(shard[:rows, :n].to(out.device)) + full = torch.cat(whole, dim=1) + for row in range(rows): + window = kv_len - next_n + row + 1 + if window <= TOPK: + count = max(window, 0) + expected[row, :count] = torch.arange(count, dtype=torch.int32, device=out.device) + continue + tokens = torch.arange(window, dtype=torch.int64, device=out.device) + composite = (ordered_keys(full[row, :window]) << _TOKEN_BITS) | tokens + top = composite.topk(TOPK).values & ((1 << _TOKEN_BITS) - 1) + expected[row] = top.sort().values.int() + return expected + + def is_valid( + self, + logits: torch.Tensor, + kv_lens: torch.Tensor, + out: torch.Tensor, + comm_buffer: torch.Tensor, + state: torch.Tensor, + next_n: int, + group_size: int, + ) -> None: + # The kernel reads kv_lens[0] only: a batch's lengths would select request 0's silently. + assert kv_lens.numel() == 1, ( + "kv_lens must hold the one generation request's indexer KV length" + ) + + +hist_topk = GlmHistTopk() diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_glm_rank_job.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_glm_rank_job.py new file mode 100644 index 000000000000..1f750a9d65b7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_glm_rank_job.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Run a collective entry's rank job from pytest, and report what it did. + +The GLM target's copy of ``tests/unittest/_torch/modeling_v2/comm/_rank_job.py``, under its own +module name so that the two never shadow each other in one pytest process: it starts the +``__op_matrix.py`` launcher beside this file. + +The entry's own launcher already owns the hard parts -- one rank per visible +device, its own process group, and a deadline it enforces by killing that +group, which is what keeps a wedged collective from taking the calling run +down with it. A collective that breaks *hangs* rather than raising, so the +deadline is load-bearing and the ``mpi_pool_executor`` fixture has none. + +So this does not reimplement any of that. It selects the devices, starts the +launcher in a fresh interpreter, and turns its exit code into an assertion. + +Fresh interpreter is required, not tidiness: the launcher owns the deadline +that kills a wedged collective's process sessions, and it must carry no CUDA +or torch.distributed state of its own -- a pytest process that has imported +``tensorrt_llm`` does. + +Everything here is started **by file path**: the ranks reach the catalog by +absolute import, so neither half needs a package context. This tree does not +have one to give: it is tests/, not a package. +""" + +from __future__ import annotations + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path + +WORLD_SIZE = 4 +"""The single-node world size the matrices run at; their 16-rank cells need a 4-node launch.""" + +_LAUNCHER_GRACE_S = 300 +"""Headroom over the entry's own deadline, so its message wins the race.""" + + +def _devices() -> str: + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible: + devices = [d for d in visible.split(",") if d.strip()] + else: + import torch + + devices = [str(i) for i in range(torch.cuda.device_count())] + assert len(devices) >= WORLD_SIZE, ( + f"this entry is certified at world size {WORLD_SIZE}; only " + f"{len(devices)} device(s) are visible" + ) + return ",".join(devices[:WORLD_SIZE]) + + +def _load_launcher_constants(launcher: Path): + """Read the entry's module-level budget without importing it by name. + + There is no dotted name to import it by -- this directory is not a + package -- and only its constants are wanted. Executing the file at module + scope is cheap and pulls in nothing but torch; the entry keeps its + tensorrt_llm imports inside the rank body for exactly that reason. + """ + spec = importlib.util.spec_from_file_location(launcher.stem, launcher) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def run(entry: str) -> None: + """Run ``__op_matrix``'s launcher over ``WORLD_SIZE`` devices.""" + launcher = Path(__file__).resolve().with_name(f"_{entry}_op_matrix.py") + env = dict(os.environ, CUDA_VISIBLE_DEVICES=_devices()) + + # The ranks import tensorrt_llm absolutely, and a source checkout is not + # necessarily installed. + # tests/unittest/_torch/modeling_v2/models/glm_moe_dsa//dsa -> repo root. + repo_root = Path(__file__).resolve().parents[8] + assert (repo_root / "tensorrt_llm").is_dir(), ( + f"expected the repo root at {repo_root}, found no tensorrt_llm/ there; " + f"this file moved without its parents[] index following" + ) + env["PYTHONPATH"] = os.pathsep.join(p for p in (str(repo_root), env.get("PYTHONPATH", "")) if p) + + # Read the budget off the entry rather than restating it: an entry that + # raises its own deadline would otherwise be killed by this one first, + # and the message a reader needs ("wedged") would be lost. + constants = _load_launcher_constants(launcher) + timeout = constants.DEADLINE_S + getattr(constants, "WEDGE_CAP_S", 0) + _LAUNCHER_GRACE_S + + completed = subprocess.run( + [sys.executable, str(launcher)], + env=env, + capture_output=True, + text=True, + timeout=timeout, + ) + if completed.returncode != 0: + raise AssertionError( + f"{launcher.name} exited {completed.returncode}\n" + f"--- stdout ---\n{completed.stdout}\n" + f"--- stderr ---\n{completed.stderr}" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_hist_topk_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_hist_topk_op_matrix.py new file mode 100644 index 000000000000..328e916d0087 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_hist_topk_op_matrix.py @@ -0,0 +1,359 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for the GLM hist_topk entry (the key-sharded decode indexer top-k). + +A collective cannot be exercised in one process, so this script is its own launcher: run plainly, +it spawns one rank per device named in CUDA_VISIBLE_DEVICES, under a deadline it enforces by +killing every rank's process group. A broken exchange waits (up to the kernel's 120 s peer timeout) +rather than raising, so the deadline is what keeps it from taking the calling run down:: + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _hist_topk_op_matrix.py + +That run drives the entry's 4-rank cells. Its 16-rank cells (c1) need 16 ranks over 4 nodes: start +this file with ``--rank-worker`` on every task, one per GPU, with ``MASTER_ADDR`` / ``MASTER_PORT`` +pointing at rank 0's node and a free port. ``RANK``, ``WORLD_SIZE``, ``LOCAL_RANK`` and +``LOCAL_WORLD_SIZE`` fall back to Slurm's ``SLURM_PROCID``, ``SLURM_NTASKS``, ``SLURM_LOCALID`` and +``SLURM_NTASKS_PER_NODE``:: + + srun -N4 --ntasks-per-node=4 env MASTER_ADDR= MASTER_PORT= \\ + python _hist_topk_op_matrix.py --rank-worker + +Every rank builds the same full logits from a seed, scores nothing itself, and hands the op its own +shard (the columns past its keys are NaN, which would outrank +inf if they were read). The checks: +the cells against ``reference`` (row sets, the order within a row being unspecified), the raw +output bit-identical on all ranks, back-to-back calls through the rotating Lamport buffers, a call +without an exchange leaving the workspace untouched, CUDA-graph replays with new logits and +lengths, and the guard. + +Not a pytest module: the ``check_*`` bodies are one fixed sequence inside a single job and share +the rank's workspace. The collected entry point is ``test_modeling_v2_hist_topk_op_matrix.py``, +which starts this file and turns its exit code into an assertion. +""" + +import os +import signal +import socket +import subprocess +import sys +import time +import zlib +from pathlib import Path +from typing import Any, List, Tuple + +import torch +import torch.distributed as dist + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 1200 + +# Bound inside the rank body, never in the launcher: importing the entry pulls in tensorrt_llm, +# which must not create CUDA or torch.distributed state in the launching process. +op: Any = None +canonical_rows: Any = None +shard_size: Any = None +shard_capacity: Any = None +validating: Any = None +WORKSPACE: Any = None +RANK = 0 +WORLD = 1 + +_SPECIAL_BITS = ( + 0x7F800000, # +inf + -0x00800000, # -inf (0xFF800000) + 0x00000000, # +0 + -0x80000000, # -0 (0x80000000) + 0x7FC00000, # NaN + -0x00400000, # NaN with the sign bit (0xFFC00000) + 0x7FFFFFFF, # NaN whose key is clamped below the empty word +) + + +def _full_logits(spec: dict, seed: int) -> torch.Tensor: + """The request's logits [rows, kv_len], the same on every rank.""" + gen = torch.Generator().manual_seed(seed) + rows, kv_len, pattern = spec["rows"], spec["kv_len"], spec["pattern"] + if pattern == "ties": + logits = torch.randint(0, 64, (rows, kv_len), generator=gen).float() + elif pattern == "equal": + logits = torch.ones(rows, kv_len) + elif pattern == "saturated": + # Sign, exponent and the top two mantissa bits equal: one top-11-bit bin holds them all. + logits = 1.0 + torch.rand(rows, kv_len, generator=gen) * 2.0**-10 + else: + logits = torch.randn(rows, kv_len, generator=gen) + if pattern == "special": + values = torch.tensor(_SPECIAL_BITS, dtype=torch.int64).to(torch.int32).view(torch.float32) + for value in values: + index = torch.randint(0, kv_len, (rows, 64), generator=gen) + logits.scatter_(1, index, value.expand(rows, 64).contiguous()) + return logits.cuda() + + +def _shard(full: torch.Tensor, kv_len: int, rank: int, width: int) -> torch.Tensor: + """Rank ``rank``'s shard logits [rows, width]: its keys, then NaN.""" + size = shard_size(kv_len, WORLD) + n = min(max(kv_len - rank * size, 0), size) + shard = torch.full((full.shape[0], width), float("nan"), device="cuda") + shard[:, :n] = full[:, rank * size : rank * size + n] + return shard + + +def _call( + full: torch.Tensor, spec: dict, width: int = 0 +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """One call of the op on this rank's shard: (out, logits, kv_lens).""" + kv_len = spec["kv_len"] + logits = _shard(full, kv_len, RANK, max(width, shard_size(kv_len, WORLD))) + kv_lens = torch.tensor([kv_len], dtype=torch.int32, device="cuda") + out = torch.full((spec["rows"], 2048), -7, dtype=torch.int32, device="cuda") + with validating(op): + op( + logits, + kv_lens, + out, + WORKSPACE.comm_buffer, + WORKSPACE.state, + spec["next_n"], + WORLD, + ) + torch.cuda.synchronize() + return out, logits, kv_lens + + +def _verify(full: torch.Tensor, spec: dict, out: torch.Tensor, kv_lens: torch.Tensor, where: str): + """The row sets against ``reference``; the raw rows bit-identical on every rank.""" + width = shard_size(spec["kv_len"], WORLD) + shards: List[torch.Tensor] = [_shard(full, spec["kv_len"], r, width) for r in range(WORLD)] + expected = op.reference( + shards, kv_lens, out, WORKSPACE.comm_buffer, WORKSPACE.state, spec["next_n"], WORLD + ) + got = canonical_rows(out) + if not torch.equal(got, expected): + bad = (got != expected).any(dim=1).nonzero().flatten().tolist() + raise AssertionError(f"{where}: rows {bad} hold another set than the reference") + gathered = [torch.empty_like(out) for _ in range(WORLD)] + dist.all_gather(gathered, out) + for rank, rows in enumerate(gathered): + assert torch.equal(rows, out), f"{where}: rank {rank} wrote other bits than rank {RANK}" + assert int(WORKSPACE.state[2]) == 0, f"{where}: the op recorded a peer timeout" + + +def check_certified_cells() -> None: + ran, skipped = 0, [] + for cell in op.CELLS: + spec = cell.spec + if spec["group_size"] != WORLD: + skipped.append(cell.why) + continue + full = _full_logits(spec, zlib.crc32(cell.why.encode()) % 2**31) + out, _, kv_lens = _call(full, spec) + _verify(full, spec, out, kv_lens, cell.why) + ran += 1 + assert ran > 0, f"no cell of hist_topk.CELLS runs at {WORLD} ranks" + if RANK == 0: + print(f"[rank 0] {ran} cells at {WORLD} ranks; not run here: {skipped}", flush=True) + + +def check_back_to_back_calls_rotate_the_lamport_buffers() -> None: + """Calls with and without an exchange, of every size, one after the other: each exchange + launch takes the next of the three buffers and clears what the previous one wrote.""" + lengths = (150_000, 1_000, 90_000, 2_051, 4 * 98_304, 60_000, 30_000, 250_000, 1_500, 5_000) + patterns = ("randn", "randn", "ties", "randn", "saturated", "equal", "special", "randn") + for step, kv_len in enumerate(lengths): + spec = dict(kv_len=kv_len, pattern=patterns[step % len(patterns)], rows=8, next_n=8) + full = _full_logits(spec, 1_000 + step) + out, _, kv_lens = _call(full, spec) + _verify(full, spec, out, kv_lens, f"call {step} ({kv_len} keys)") + + +def check_rows_without_an_exchange_leave_the_workspace_alone() -> None: + state = WORKSPACE.state.clone() + buffer = WORKSPACE.comm_buffer.clone() + spec = dict(kv_len=2_000, pattern="randn", rows=8, next_n=8) + full = _full_logits(spec, 2_000) + out, _, kv_lens = _call(full, spec) + _verify(full, spec, out, kv_lens, "2,000 keys") + assert torch.equal(WORKSPACE.state, state), "a launch without an exchange moved the state" + assert torch.equal(WORKSPACE.comm_buffer, buffer), "a launch without an exchange wrote" + + +def check_cuda_graph_replay_follows_new_logits_and_lengths() -> None: + """The decode step is one captured graph: replays read the logits and the KV length in place, + and every replay is an exchange launch of its own.""" + width = shard_capacity(300_000, WORLD) + first = dict(kv_len=200_000, pattern="randn", rows=8, next_n=8) + full = _full_logits(first, 3_000) + out, logits, kv_lens = _call(full, first, width) + _verify(full, first, out, kv_lens, "graph warm-up") + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(logits, kv_lens, out, WORKSPACE.comm_buffer, WORKSPACE.state, 8, WORLD) + for step, kv_len in enumerate((180_000, 300_000, 1_900, 120_001)): + spec = dict(kv_len=kv_len, pattern="ties" if step % 2 else "randn", rows=8, next_n=8) + full = _full_logits(spec, 3_001 + step) + logits.copy_(_shard(full, kv_len, RANK, width)) + kv_lens.fill_(kv_len) + out.fill_(-7) + dist.barrier() + graph.replay() + torch.cuda.synchronize() + _verify(full, spec, out, kv_lens, f"replay {step} ({kv_len} keys)") + + +def check_guard_refuses_a_batch_of_kv_lens() -> None: + spec = dict(kv_len=100_000, pattern="randn", rows=8, next_n=8) + logits = _shard(_full_logits(spec, 4_000), spec["kv_len"], RANK, shard_size(100_000, WORLD)) + kv_lens = torch.tensor([100_000, 7], dtype=torch.int32, device="cuda") + out = torch.empty(8, 2048, dtype=torch.int32, device="cuda") + try: + with validating(op): + op(logits, kv_lens, out, WORKSPACE.comm_buffer, WORKSPACE.state, 8, WORLD) + except AssertionError as error: + assert "kv_lens" in str(error), error + else: + raise AssertionError("is_valid admitted the KV lengths of two requests") + + +CHECKS = ( + check_certified_cells, + check_back_to_back_calls_rotate_the_lamport_buffers, + check_rows_without_an_exchange_leave_the_workspace_alone, + check_cuda_graph_replay_follows_new_logits_and_lengths, + check_guard_refuses_a_batch_of_kv_lens, +) + + +def _env_int(name: str, slurm_name: str, default: int) -> int: + value = os.environ.get(name) or os.environ.get(slurm_name) + return int(value) if value else default + + +def _run_one_rank() -> int: + """Body of one rank: run every check; a failing rank exits nonzero and the launcher kills + its siblings.""" + global op, canonical_rows, shard_size, shard_capacity, validating, WORKSPACE, RANK, WORLD + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = _env_int("RANK", "SLURM_PROCID", 0) + world = _env_int("WORLD_SIZE", "SLURM_NTASKS", 1) + local_rank = _env_int("LOCAL_RANK", "SLURM_LOCALID", rank) + local_world = _env_int("LOCAL_WORLD_SIZE", "SLURM_NTASKS_PER_NODE", world) + torch.cuda.set_device(local_rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", + init_method="env://", + rank=rank, + world_size=world, + device_id=torch.device("cuda", local_rank), + ) + + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa import ( # noqa: E501 + _hist_topk, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa import ( # noqa: E501 + hist_topk as entry, + ) + from tensorrt_llm.mapping import Mapping + + # A spawned worker, not a collected test: the repo's __extra_import_path__ is not in play, so + # the shared helper is reached the plain way, inside this rank process only. + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating as _validating_cm + + op = entry.hist_topk + canonical_rows = entry.canonical_rows + shard_size = _hist_topk.shard_size + shard_capacity = _hist_topk.shard_capacity + validating = _validating_cm + RANK, WORLD = rank, world + assert 2 <= WORLD <= _hist_topk.MAX_RANKS, f"a {WORLD}-rank group" + mapping = Mapping(world_size=world, rank=rank, gpus_per_node=local_world, tp_size=world) + WORKSPACE = _hist_topk.create_hist_topk_workspace(mapping, comm=dist.group.WORLD) + + for check in CHECKS: + try: + check() + dist.barrier() + except BaseException: + import traceback + + print(f"[rank {RANK}] FAILED {check.__name__}", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + # Exit hard: a rank that leaves a collective early wedges its siblings, and the + # launcher kills the rest when it sees this exit. + os._exit(13) + torch.cuda.synchronize() + print(f"[rank {RANK}] {len(CHECKS)} checks passed", flush=True) + # Exit without tearing the process group down: these are throwaway processes. + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """Start one worker process per claimed device on this node; kill them all on the first + nonzero exit or at the deadline.""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the devices this run owns, e.g. 0,1,2,3" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size >= 2, f"CUDA_VISIBLE_DEVICES names {world_size} device(s); need at least 2" + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + LOCAL_WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + ) + # Drop the launching step's PMIx/PMI/OMPI/Slurm namespace from the children: they would + # otherwise register as participants of that step, and the first hard exit would have the + # step killed before anything is reported. + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + print(f"[launcher] spawning {world_size} ranks, rendezvous on :{port}", flush=True) + children = [ + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + for r in range(world_size) + ] + + def _kill_all() -> None: + for p in children: + try: + os.killpg(p.pid, signal.SIGKILL) + except ProcessLookupError: + pass + for p in children: + p.wait() + + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" + break + for r in sorted(live): + code = children[r].poll() + if code is None: + continue + live.discard(r) + if code != 0: + failure = f"rank {r} exited {code}" + break + time.sleep(0.5) + if failure is not None: + _kill_all() + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + + +if __name__ == "__main__": + if _WORKER_FLAG in sys.argv: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_shard_meta.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_shard_meta.py new file mode 100644 index 000000000000..ff3d0961e9ec --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_shard_meta.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM hist_shard_meta entry: integer outputs, compared exactly. + +The request's INDEX_KEY block-table row holds distinct random page ids (with slack past the +request's pages, as the engine's table rows have), and both outputs start out as a sentinel so a +missing write shows. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa._hist_topk import ( # noqa: E501 + shard_size, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.hist_shard_meta import ( # noqa: E501 + hist_shard_meta as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +assert torch.cuda.is_available(), "hist_shard_meta requires a CUDA device" + + +def _build(spec, seed: int) -> dict: + gen = torch.Generator().manual_seed(seed) + tpb = spec["tokens_per_block"] + pages = -(-spec["kv_len"] // tpb) + 7 + row = torch.randperm(4 * pages, generator=gen)[:pages].to(torch.int32) + width = -(-spec["capacity"] // tpb) + return dict( + kv_lens=torch.tensor([spec["kv_len"]], dtype=torch.int32, device="cuda"), + block_table=row.cuda(), + shard_ctx=torch.full((1,), -7, dtype=torch.int32, device="cuda"), + shard_block_table=torch.full((1, width), -7, dtype=torch.int32, device="cuda"), + group_size=spec["group_size"], + rank=spec["rank"], + tokens_per_block=tpb, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + ref_ctx, ref_table = op.reference(**args) + ctx, table = op(**args) + torch.cuda.synchronize() + assert ctx is args["shard_ctx"] and table is args["shard_block_table"] + assert torch.equal(ctx, ref_ctx), (ctx.tolist(), ref_ctx.tolist()) + assert torch.equal(table, ref_table) + + +def test_shards_of_all_ranks_tile_the_request(): + """Independent of the reference: the 16 shards' pages, in rank order, are the request's + pages, and their key counts sum to the request's length.""" + spec = op.CELLS[0].spec + args = _build(spec, seed=61) + pages, keys = [], 0 + size = shard_size(spec["kv_len"], spec["group_size"]) + for rank in range(spec["group_size"]): + ctx, table = op(**dict(args, rank=rank)) + n = int(ctx[0]) + held = min(max(spec["kv_len"] - rank * size, 0), size) + assert n == max(held, 1) + keys += held + pages += table[0, : -(-held // spec["tokens_per_block"])].tolist() + assert keys == spec["kv_len"] + total_pages = -(-spec["kv_len"] // spec["tokens_per_block"]) + assert pages == args["block_table"][:total_pages].tolist() + + +def test_guard_refuses_a_two_dimensional_table(): + args = _build(op.CELLS[0].spec, seed=62) + with pytest.raises(AssertionError, match="one-dimensional"): + with validating(op): + op(**dict(args, block_table=args["block_table"][None, :])) + + +def test_guard_refuses_a_shard_table_narrower_than_the_shard(): + args = _build(op.CELLS[0].spec, seed=63) + with pytest.raises(AssertionError, match="shard_block_table holds"): + with validating(op): + op(**dict(args, shard_block_table=args["shard_block_table"][:, :100].contiguous())) + + +def test_guard_refuses_a_row_shorter_than_the_shard(): + args = _build(op.CELLS[1].spec, seed=64) + with pytest.raises(AssertionError, match="run past"): + with validating(op): + op(**dict(args, block_table=args["block_table"][:1000].contiguous())) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_topk_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_topk_op_matrix.py new file mode 100644 index 000000000000..d3ac1acda903 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_topk_op_matrix.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the GLM hist_topk entry's certification matrix. + +The matrix itself is ``_hist_topk_op_matrix.py`` beside this file, its own 4-rank launcher (see +``_glm_rank_job`` for why it is left intact). It drives the entry's 4-rank cells; the 16-rank cells +need the 4-node launch its docstring describes. +""" + +import _glm_rank_job +import pytest +import torch + +from tensorrt_llm._utils import get_sm_version + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +assert torch.cuda.is_available(), "hist_topk requires CUDA devices" + + +# Each case spawns its own 4 rank processes over the visible devices; under xdist several +# workers would fight for the same GPUs. +@pytest.mark.no_xdist +def test_hist_topk_op_matrix() -> None: + _glm_rank_job.run("hist_topk") From cdd392bcfc39c02dcbfc8630c787ec4284f77fb9 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:19:48 -0700 Subject: [PATCH 31/53] Add the GLM fused DSA indexer prologue as target-local Triton One Triton kernel computes the indexer prologue of a decode step: the k_norm LayerNorm, RoPE of the Q and K rope slices, E4M3 Q and K with one UE8M0 (or amax / 448) scale per 128 values, the weights scaled by the Q scales and softmax_scale * n_heads**-0.5, and the K rows written into the indexer K cache at the step's slot mappings, in the cache layout and slot rule of indexer_k_cache_scatter_op. The arithmetic rounds where the per-kernel chain rounds. With PDL the kernel can trigger its dependent right after its grid-dependency wait and load the step inputs before it. The tp16 target gets two entries: indexer_k_tail, the K part alone (the Q part compiled out) for a step whose wq_b epilogue computes the Q part, and indexer_prologue, the whole prologue. Both are certified on sm_107 against plain-torch references, with the cache bytes checked exactly. Signed-off-by: Dongfeng Yu --- .../ops/dsa/_indexer_prologue.py | 470 ++++++++++++++++++ .../ops/dsa/_indexer_reference.py | 149 ++++++ .../ops/dsa/indexer_k_tail.py | 277 +++++++++++ .../ops/dsa/indexer_prologue.py | 272 ++++++++++ .../dsa/_indexer_inputs.py | 103 ++++ .../dsa/test_modeling_v2_indexer_k_tail.py | 191 +++++++ .../dsa/test_modeling_v2_indexer_prologue.py | 172 +++++++ 7 files changed, 1634 insertions(+) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_indexer_prologue.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_indexer_reference.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/indexer_k_tail.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/indexer_prologue.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_indexer_inputs.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_indexer_k_tail.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_indexer_prologue.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_indexer_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_indexer_prologue.py new file mode 100644 index 000000000000..5ecdcf2bc7f0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_indexer_prologue.py @@ -0,0 +1,470 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused DSA indexer prologue: one Triton launch for the indexer K and Q of a decode step. + +``_dsa_indexer_prologue_fused_kernel`` computes in one launch what the indexer's per-kernel chain +computes between the indexer projections and the paged MQA logits: + + 1. ``k_norm`` LayerNorm over the indexer head dim; + 2. RoPE on the Q and K rope slices (the flashinfer ``apply_rope_with_cos_sin_cache_inplace`` + expression); + 3. ``[pe | nope]`` FP8 E4M3 quantization of Q and K with one scale per 128 values + (``fused_cat_fp8``'s rule); + 4. ``weights * q_scale * softmax_scale * n_heads**-0.5``; + 5. the scatter of each K row (128 FP8 bytes and a 4-byte scale) into the indexer K cache + (``indexer_k_cache_scatter_op``'s layout and slot rule). + +Each step has the numerics of the chain's kernel: fp32 math on bf16-rounded intermediates (the +LayerNorm output and the RoPE output are rounded to bf16 where the chain materialises bf16 +tensors), ``fused_cat_fp8``'s scale (UE8M0 power-of-two round-up or plain amax / 448), its +``rcp.approx.ftz`` reciprocal and RN-satfinite E4M3 conversion, and the scatter kernel's flat-index +to strided-offset unravel. + +``launch_indexer_prologue`` runs the whole prologue. ``launch_indexer_k_tail`` runs the same kernel +with the Q part compiled out (``WRITE_Q=False``), one program per token: the K part of a step whose +indexer wq_b epilogue computes the Q part. + +With ``use_pdl`` the kernel is launched with programmatic dependent launch and waits on +``griddepcontrol.wait`` before it reads its inputs: + +- ``launch_dependents``: each program triggers the dependent launch right after its wait, so the + next kernel on the stream can start while this one runs. Only for a dependent that reads this + kernel's outputs (Q, the weights, the new K-cache rows) after its own wait and writes nothing + before it: the early-KV decode MQA-logits kernel. +- ``prewait``: each program loads the position, the cos/sin row, the K LayerNorm weight and bias + and the slot mappings before its wait, and Q, K and the weights right after it. Only when the + positions and slot mappings were written before the kernels this one depends on started (the + step's metadata preparation). +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _rcp_approx_ftz(x): + # fused_cat_fp8 uses the MUFU reciprocal (rcp.approx.ftz.f32), not IEEE 1/x. + return tl.inline_asm_elementwise( + "rcp.approx.ftz.f32 $0, $1;", + "=r,r", + [x], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _fp8_row_scale(amax, UE8M0: tl.constexpr): + """fused_cat_fp8 per-row scale: amax clamped at 1e-12, / 448, optional UE8M0 round-up.""" + amax = tl.maximum(amax, 1.0e-12) + ratio = amax * 0.002232142857142857 # 1.0f / 448.0f + if UE8M0: + bits = ratio.to(tl.uint32, bitcast=True) + mant = bits & 0x007FFFFF + exp_bits = bits & 0x7F800000 + exp_bits = tl.where(mant != 0, exp_bits + 0x00800000, exp_bits) + scale = exp_bits.to(tl.float32, bitcast=True) + else: + scale = ratio + return scale + + +@triton.jit +def _rope_terms(d, ROPE_DIM: tl.constexpr, INTERLEAVE: tl.constexpr): + """Partner column, cos/sin column and sign for column ``d`` (flashinfer layout).""" + HALF: tl.constexpr = ROPE_DIM // 2 + is_rope = d < ROPE_DIM + if INTERLEAVE: + # GPT-J style: pairs (2i, 2i + 1) share cos/sin column i; even lanes subtract. + partner = tl.where(is_rope, d ^ 1, d) + cs_col = d // 2 + neg = (d % 2) == 0 + else: + # NeoX style: halves [0, HALF) and [HALF, ROPE_DIM); the first half subtracts. + partner = tl.where(d < HALF, d + HALF, tl.where(is_rope, d - HALF, d)) + cs_col = d % HALF + neg = d < HALF + cs_col = tl.where(is_rope, cs_col, 0) + return is_rope, partner, cs_col, neg + + +@triton.jit +def _apply_rope(x, xp, c, s, neg): + # flashinfer: vec[i] = vec[i] * cos[i] + (neg ? -perm[i] : perm[i]) * sin[i]; + # nvcc contracts this to fma(vec, cos, (+-perm) * sin). + t = tl.where(neg, -xp, xp) * s + return tl.fma(x, c, t) + + +@triton.jit +def _unravel_offset(flat, d1, d2, d3, s0, s1, s2, s3): + """indexerKCacheScatter's flatIndexToMemoryOffset over a [d0, d1, d2, d3] strided view.""" + i3 = flat % d3 + flat = flat // d3 + i2 = flat % d2 + flat = flat // d2 + i1 = flat % d1 + i0 = flat // d1 + return i0 * s0 + i1 * s1 + i2 * s2 + i3 * s3 + + +@triton.jit +def _dsa_indexer_prologue_fused_kernel( + # inputs + k_ptr, # [N, D] pre-norm indexer K (bf16 or fp32; rounded to bf16 first) + k_stride, + ln_w_ptr, # [D] LayerNorm weight + ln_b_ptr, # [D] LayerNorm bias + q_ptr, # [N, H, D] bf16 (wq_b output view, unit inner stride) + q_stride_tok, + q_stride_head, + w_ptr, # [N, H] weights_proj output (bf16 or fp32) + w_stride, + pos_ptr, # [N] positions + cos_sin_ptr, # [max_pos, ROPE_DIM] fp32: [cos(0..HALF) | sin(0..HALF)] + cos_sin_stride, + # outputs + q_fp8_ptr, # [N, H, D] fp8 e4m3 (contiguous) + q_scale_ptr, # [N, H] fp32 + k_fp8_ptr, # [N, D] fp8 e4m3 (contiguous) + k_scale_ptr, # [N] fp32 + w_out_ptr, # [N, H] fp32 + # indexer K cache (uint8 [d0, d1, d2, d3] strided view) + slot mappings + k_cache_ptr, + slot_fp8_ptr, + slot_scale_ptr, + n_write, + cache_d1, + cache_d2, + cache_d3, + cache_s0, + cache_s1, + cache_s2, + cache_s3, + # scalars + eps, + weight_scale, + H: tl.constexpr, + D: tl.constexpr, + ROPE_DIM: tl.constexpr, + HEADS_PER_PROG: tl.constexpr, + INTERLEAVE: tl.constexpr, + UE8M0: tl.constexpr, + WRITE_CACHE: tl.constexpr, + USE_PDL: tl.constexpr, + LAUNCH_DEPENDENTS: tl.constexpr, + PREWAIT: tl.constexpr, + WRITE_Q: tl.constexpr, +): + tok = tl.program_id(0) + hblk = tl.program_id(1) + HALF: tl.constexpr = ROPE_DIM // 2 + if PREWAIT: + # Positions, cos/sin, the K LayerNorm and the slot mappings: written before the + # kernels this one depends on run. + tok64 = tok.to(tl.int64) + d = tl.arange(0, D) + is_rope, partner, cs_col, neg = _rope_terms(d, ROPE_DIM, INTERLEAVE) + k_prog = hblk == 0 + pos = tl.load(pos_ptr + tok).to(tl.int64) + cs_row = cos_sin_ptr + pos * cos_sin_stride + cos_v = tl.load(cs_row + cs_col, mask=is_rope, other=1.0) + sin_v = tl.load(cs_row + HALF + cs_col, mask=is_rope, other=0.0) + lw = tl.load(ln_w_ptr + d, mask=k_prog, other=0.0).to(tl.float32) + lb = tl.load(ln_b_ptr + d, mask=k_prog, other=0.0).to(tl.float32) + lwp = tl.load(ln_w_ptr + partner, mask=k_prog, other=0.0).to(tl.float32) + lbp = tl.load(ln_b_ptr + partner, mask=k_prog, other=0.0).to(tl.float32) + if WRITE_CACHE: + slot = tl.load(slot_fp8_ptr + tok, mask=k_prog & (tok < n_write), other=-1) + slot_s = tl.load(slot_scale_ptr + tok, mask=k_prog & (tok < n_write), other=-1) + if USE_PDL: + tl.extra.cuda.gdc_wait() + if LAUNCH_DEPENDENTS: + tl.extra.cuda.gdc_launch_dependents() + + if not PREWAIT: + tok64 = tok.to(tl.int64) + d = tl.arange(0, D) + is_rope, partner, cs_col, neg = _rope_terms(d, ROPE_DIM, INTERLEAVE) + + pos = tl.load(pos_ptr + tok).to(tl.int64) + cs_row = cos_sin_ptr + pos * cos_sin_stride + cos_v = tl.load(cs_row + cs_col, mask=is_rope, other=1.0) + sin_v = tl.load(cs_row + HALF + cs_col, mask=is_rope, other=0.0) + + if WRITE_Q: + # ---------------- Q: HEADS_PER_PROG heads of this token ---------------- + h = hblk * HEADS_PER_PROG + tl.arange(0, HEADS_PER_PROG) + q_row = q_ptr + tok64 * q_stride_tok + h[:, None].to(tl.int64) * q_stride_head + qx = tl.load(q_row + d[None, :]).to(tl.float32) + qp = tl.load(q_row + partner[None, :]).to(tl.float32) + if PREWAIT: + # Every load from the kernels this one waits for, before the first store. + w = tl.load(w_ptr + tok64 * w_stride + h).to(tl.float32) + k_row = k_ptr + tok64 * k_stride + kx = tl.load(k_row + d, mask=k_prog, other=0.0).to(tl.bfloat16).to(tl.float32) + kxp = tl.load(k_row + partner, mask=k_prog, other=0.0).to(tl.bfloat16).to(tl.float32) + q_rot = _apply_rope(qx, qp, cos_v[None, :], sin_v[None, :], neg[None, :]) + qv = tl.where(is_rope[None, :], q_rot, qx).to(tl.bfloat16).to(tl.float32) + q_amax = tl.max(tl.abs(qv), axis=1) + q_scale = _fp8_row_scale(q_amax, UE8M0) + q_inv = _rcp_approx_ftz(q_scale) + q8 = (qv * q_inv[:, None]).to(tl.float8e4nv) + qo = (tok64 * H + h[:, None]) * D + d[None, :] + tl.store(q_fp8_ptr + qo, q8) + tl.store(q_scale_ptr + tok64 * H + h, q_scale) + + # weights * q_scale * (softmax_scale * n_heads^-0.5), same op order as _scale(). + if not PREWAIT: + w = tl.load(w_ptr + tok64 * w_stride + h).to(tl.float32) + tl.store(w_out_ptr + tok64 * H + h, (w * q_scale) * weight_scale) + elif PREWAIT: + # K only: the K row, right after the wait. + k_row = k_ptr + tok64 * k_stride + kx = tl.load(k_row + d, mask=k_prog, other=0.0).to(tl.bfloat16).to(tl.float32) + kxp = tl.load(k_row + partner, mask=k_prog, other=0.0).to(tl.bfloat16).to(tl.float32) + + # ---------------- K: once per token ---------------- + if hblk == 0: + if not PREWAIT: + k_row = k_ptr + tok64 * k_stride + kx = tl.load(k_row + d).to(tl.bfloat16).to(tl.float32) + kxp = tl.load(k_row + partner).to(tl.bfloat16).to(tl.float32) + mean = tl.sum(kx, axis=0) / D + kc = kx - mean + var = tl.sum(kc * kc, axis=0) / D + rstd = tl.math.rsqrt(var + eps) + if not PREWAIT: + lw = tl.load(ln_w_ptr + d).to(tl.float32) + lb = tl.load(ln_b_ptr + d).to(tl.float32) + lwp = tl.load(ln_w_ptr + partner).to(tl.float32) + lbp = tl.load(ln_b_ptr + partner).to(tl.float32) + kn = (kc * rstd * lw + lb).to(tl.bfloat16).to(tl.float32) + knp = ((kxp - mean) * rstd * lwp + lbp).to(tl.bfloat16).to(tl.float32) + k_rot = _apply_rope(kn, knp, cos_v, sin_v, neg) + kv = tl.where(is_rope, k_rot, kn).to(tl.bfloat16).to(tl.float32) + k_amax = tl.max(tl.abs(kv), axis=0) + k_scale = _fp8_row_scale(k_amax, UE8M0) + # inline asm needs a block operand: broadcast the row scale over D. + k_inv = _rcp_approx_ftz(k_scale + tl.zeros([D], dtype=tl.float32)) + k8 = (kv * k_inv).to(tl.float8e4nv) + tl.store(k_fp8_ptr + tok64 * D + d, k8) + tl.store(k_scale_ptr + tok64, k_scale) + + if WRITE_CACHE: + if tok < n_write: + if not PREWAIT: + slot = tl.load(slot_fp8_ptr + tok) + slot_s = tl.load(slot_scale_ptr + tok) + if (slot >= 0) & (slot_s >= 0): + # 4-byte chunks, each chunk base unravelled (as the scatter kernel). + chunk = slot + (d // 4) * 4 + dst = _unravel_offset( + chunk, cache_d1, cache_d2, cache_d3, cache_s0, cache_s1, cache_s2, cache_s3 + ) + (d % 4) + tl.store(k_cache_ptr + dst, k8.to(tl.uint8, bitcast=True)) + b = tl.arange(0, 4) + sbits = k_scale.to(tl.int32, bitcast=True) + sbytes = ((sbits >> (b * 8)) & 0xFF).to(tl.uint8) + dst_s = _unravel_offset( + slot_s, cache_d1, cache_d2, cache_d3, cache_s0, cache_s1, cache_s2, cache_s3 + ) + tl.store(k_cache_ptr + dst_s + b, sbytes) + + +def _heads_per_prog(n_heads: int) -> int: + for hp in (8, 4, 2, 1): + if n_heads % hp == 0: + return hp + return 1 + + +def _cache_args( + k_cache: torch.Tensor, + slot_mapping_fp8: torch.Tensor, + slot_mapping_scale: torch.Tensor, + num_cache_tokens: int, +) -> tuple: + """The K-cache operands: the uint8 [d0, d1, d2, d3] view, the slot mappings, the row count, + and the view's inner sizes and strides.""" + cache_u8 = k_cache.view(torch.uint8) + cd = cache_u8.shape + cs = cache_u8.stride() + return ( + cache_u8, + slot_mapping_fp8, + slot_mapping_scale, + int(num_cache_tokens), + cd[1], + cd[2], + cd[3], + cs[0], + cs[1], + cs[2], + cs[3], + ) + + +def launch_indexer_prologue( + indexer_k: torch.Tensor, + weights: torch.Tensor, + q: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + eps: float, + weight_scale: float, + rope_dim: int, + interleave: bool, + use_ue8m0: bool, + k_cache: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + num_cache_tokens: int, + use_pdl: bool, + launch_dependents: bool, + prewait: bool, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """The whole prologue in one launch. + + Args: + indexer_k: [N, D] pre-norm indexer K (bf16, or fp32 which is rounded to bf16), unit inner + stride. + weights: [N, H] weights_proj output (bf16 or fp32), unit inner stride. + q: [N, H, D] bf16 indexer Q (the wq_b output), unit inner stride. + positions: [N] (any int dtype) rope positions. + cos_sin: [max_pos, rope_dim] fp32 cos/sin cache, ``[cos | sin]`` per row. + k_cache / slot_mapping_*: when ``k_cache`` is given and ``num_cache_tokens`` > 0, the first + ``num_cache_tokens`` K rows are also written into the indexer K cache (the layout and + slot rule of ``indexer_k_cache_scatter_op``). + use_pdl / launch_dependents / prewait: see the module docstring. + + Returns (q_fp8 [N, H, D], k_fp8 [N, D], k_scale [N, 1], weights [N, H] fp32, + q_scale [N, H, 1]). + """ + n_tok, n_heads, head_dim = q.shape + dev = q.device + q_fp8 = torch.empty((n_tok, n_heads, head_dim), dtype=torch.float8_e4m3fn, device=dev) + q_scale = torch.empty((n_tok, n_heads, 1), dtype=torch.float32, device=dev) + k_fp8 = torch.empty((n_tok, head_dim), dtype=torch.float8_e4m3fn, device=dev) + k_scale = torch.empty((n_tok, 1), dtype=torch.float32, device=dev) + w_out = torch.empty((n_tok, n_heads), dtype=torch.float32, device=dev) + if n_tok == 0: + return q_fp8, k_fp8, k_scale, w_out, q_scale + + write_cache = k_cache is not None and num_cache_tokens > 0 + if write_cache: + cache_args = _cache_args(k_cache, slot_mapping_fp8, slot_mapping_scale, num_cache_tokens) + else: + cache_args = (k_fp8, k_scale, k_scale, 0, 1, 1, 1, 0, 0, 0, 0) + + hp = _heads_per_prog(n_heads) + _dsa_indexer_prologue_fused_kernel[(n_tok, n_heads // hp)]( + indexer_k, + indexer_k.stride(0), + ln_weight, + ln_bias, + q, + q.stride(0), + q.stride(1), + weights, + weights.stride(0), + positions, + cos_sin, + cos_sin.stride(0), + q_fp8, + q_scale, + k_fp8, + k_scale, + w_out, + *cache_args, + float(eps), + float(weight_scale), + H=n_heads, + D=head_dim, + ROPE_DIM=rope_dim, + HEADS_PER_PROG=hp, + INTERLEAVE=bool(interleave), + UE8M0=bool(use_ue8m0), + WRITE_CACHE=write_cache, + USE_PDL=use_pdl, + LAUNCH_DEPENDENTS=use_pdl and launch_dependents, + PREWAIT=use_pdl and prewait, + WRITE_Q=True, + num_warps=4, + launch_pdl=use_pdl, + ) + return q_fp8, k_fp8, k_scale, w_out, q_scale + + +def launch_indexer_k_tail( + indexer_k: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + eps: float, + rope_dim: int, + interleave: bool, + use_ue8m0: bool, + k_fp8: torch.Tensor, + k_scale: torch.Tensor, + k_cache: torch.Tensor, + slot_mapping_fp8: torch.Tensor, + slot_mapping_scale: torch.Tensor, + num_cache_tokens: int, + use_pdl: bool, + launch_dependents: bool, + prewait: bool, +) -> None: + """The K part of ``launch_indexer_prologue``, one program per token: the K LayerNorm, RoPE and + FP8 quantization of K into ``k_fp8`` [N, D] / ``k_scale`` [N, 1], and the first + ``num_cache_tokens`` rows into the indexer K cache. The kernel and its K arithmetic are those of + the whole prologue; only its Q part is compiled out (H and HEADS_PER_PROG select nothing there). + """ + n_tok = indexer_k.shape[0] + if n_tok == 0: + return + _dsa_indexer_prologue_fused_kernel[(n_tok, 1)]( + indexer_k, + indexer_k.stride(0), + ln_weight, + ln_bias, + # The Q operands are not read without WRITE_Q. + indexer_k, + 0, + 0, + indexer_k, + 0, + positions, + cos_sin, + cos_sin.stride(0), + k_fp8, + k_scale, + k_fp8, + k_scale, + k_scale, + *_cache_args(k_cache, slot_mapping_fp8, slot_mapping_scale, num_cache_tokens), + float(eps), + 1.0, + H=1, + D=indexer_k.shape[1], + ROPE_DIM=rope_dim, + HEADS_PER_PROG=1, + INTERLEAVE=bool(interleave), + UE8M0=bool(use_ue8m0), + WRITE_CACHE=True, + USE_PDL=use_pdl, + LAUNCH_DEPENDENTS=use_pdl and launch_dependents, + PREWAIT=use_pdl and prewait, + WRITE_Q=False, + num_warps=4, + launch_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_indexer_reference.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_indexer_reference.py new file mode 100644 index 000000000000..a35396b535d0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_indexer_reference.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Plain-torch references of the indexer prologue's arithmetic, shared by its two entries. + +fp32 math rounded where the kernel rounds: the LayerNorm output and the RoPE output to bf16, one +rounding for each fused multiply-add, ``fused_cat_fp8``'s per-row scale and E4M3 conversion, and +the indexer K-cache scatter's byte placement. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch + + +def fma(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> torch.Tensor: + """fp32 ``a * b + c`` with one rounding: the fp64 product of two fp32 values is exact.""" + return (a.double() * b.double() + c.double()).float() + + +def rope_rows( + x: torch.Tensor, positions: torch.Tensor, cos_sin: torch.Tensor, rope_dim: int, interleave: bool +) -> torch.Tensor: + """RoPE on the first ``rope_dim`` values of the fp32 rows ``x`` [..., D] at ``positions`` + (one per leading row), rounded to bf16 and returned as fp32. + + ``cos_sin`` rows are ``[cos | sin]`` of ``rope_dim / 2`` values each. Interleaved pairs + (2i, 2i + 1) share column i and the even value subtracts; otherwise the two halves pair up and + the first half subtracts. + """ + d = torch.arange(x.shape[-1], device=x.device) + half = rope_dim // 2 + is_rope = d < rope_dim + if interleave: + partner = torch.where(is_rope, d ^ 1, d) + cs_col = d // 2 + neg = d % 2 == 0 + else: + partner = torch.where(d < half, d + half, torch.where(is_rope, d - half, d)) + cs_col = d % half + neg = d < half + cs_col = torch.where(is_rope, cs_col, torch.zeros_like(cs_col)) + table = cos_sin[positions.long()] + while table.dim() < x.dim(): + table = table.unsqueeze(-2) + cos = torch.where(is_rope, table[..., cs_col], torch.ones_like(table[..., cs_col])) + sin = torch.where(is_rope, table[..., half + cs_col], torch.zeros_like(table[..., cs_col])) + xp = x[..., partner] + rotated = fma(x, cos, torch.where(neg, -xp, xp) * sin) + return torch.where(is_rope, rotated, x).to(torch.bfloat16).float() + + +def quantize_rows(v: torch.Tensor, use_ue8m0: bool) -> Tuple[torch.Tensor, torch.Tensor]: + """``fused_cat_fp8``'s rule over the last dim of fp32 ``v``: (E4M3 values, fp32 scales with a + trailing dim of 1). scale = max(amax, 1e-12) / 448, rounded up to a power of two with + ``use_ue8m0``; values = e4m3(v * (1 / scale)).""" + amax = v.abs().amax(dim=-1, keepdim=True).clamp_min(1.0e-12) + ratio = amax * 0.002232142857142857 + if use_ue8m0: + bits = ratio.view(torch.int32) + exp_bits = bits & 0x7F800000 + exp_bits = torch.where((bits & 0x007FFFFF) != 0, exp_bits + 0x00800000, exp_bits) + scale = exp_bits.view(torch.float32) + else: + scale = ratio + # The kernel multiplies by the MUFU reciprocal, exact for a power-of-two scale. + return (v * (1.0 / scale)).to(torch.float8_e4m3fn), scale + + +def reference_k_rows( + indexer_k: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + eps: float, + rope_dim: int, + interleave: bool, + use_ue8m0: bool, +) -> Tuple[torch.Tensor, torch.Tensor]: + """``(k_fp8 [N, D], k_scale [N, 1])``: LayerNorm of bf16(indexer_k) rounded to bf16, RoPE + rounded to bf16, then the E4M3 rule.""" + rows = indexer_k.shape[0] + kx = indexer_k.to(torch.bfloat16).float() + mean = kx.mean(dim=1, keepdim=True) + kc = kx - mean + rstd = torch.rsqrt((kc * kc).mean(dim=1, keepdim=True) + eps) + kn = fma(kc * rstd, ln_weight.float()[None, :], ln_bias.float()[None, :]).to(torch.bfloat16) + kv = rope_rows(kn.float(), positions[:rows], cos_sin, rope_dim, interleave) + return quantize_rows(kv, use_ue8m0) + + +def reference_q_rows( + q: torch.Tensor, + weights: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + weight_scale: float, + rope_dim: int, + interleave: bool, + use_ue8m0: bool, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``(q_fp8 [N, H, D], q_scale [N, H, 1], weights [N, H] fp32)``: RoPE of the bf16 Q rounded + to bf16, the E4M3 rule per head, and ``(weights * q_scale) * weight_scale`` in fp32.""" + rows = q.shape[0] + qv = rope_rows(q.float(), positions[:rows], cos_sin, rope_dim, interleave) + q8, q_scale = quantize_rows(qv, use_ue8m0) + w = (weights.float() * q_scale[..., 0]) * torch.tensor(weight_scale, dtype=torch.float32) + return q8, q_scale, w + + +def unravel(flat: torch.Tensor, shape: torch.Size) -> Tuple[torch.Tensor, ...]: + """The 4-D index of each logical flat index of a [d0, d1, d2, d3] view.""" + d1, d2, d3 = shape[1:] + i3 = flat % d3 + flat = flat // d3 + i2 = flat % d2 + flat = flat // d2 + i1 = flat % d1 + return flat // d1, i1, i2, i3 + + +def write_cache_rows( + cache: torch.Tensor, + k8: torch.Tensor, + k_scale: torch.Tensor, + slot_mapping_fp8: torch.Tensor, + slot_mapping_scale: torch.Tensor, + num_cache_tokens: int, +) -> None: + """The indexer K-cache scatter, in place: for each of the first ``num_cache_tokens`` rows with + both slots >= 0, its 128 FP8 bytes at the logical flat byte index ``slot_mapping_fp8[row]`` of + the 4-D byte view ``cache`` and its fp32 scale's 4 little-endian bytes at + ``slot_mapping_scale[row]``.""" + view = cache.view(torch.uint8) + fp8_slots = slot_mapping_fp8[:num_cache_tokens].long() + scale_slots = slot_mapping_scale[:num_cache_tokens].long() + keep = (fp8_slots >= 0) & (scale_slots >= 0) + rows = torch.nonzero(keep).flatten() + if rows.numel() == 0: + return + data = k8.view(torch.uint8)[rows] + scale_bytes = k_scale.reshape(-1, 1).contiguous().view(torch.uint8)[rows] + width = data.shape[1] + fp8_index = fp8_slots[rows, None] + torch.arange(width, device=view.device) + scale_index = scale_slots[rows, None] + torch.arange(4, device=view.device) + view[unravel(fp8_index, view.shape)] = data + view[unravel(scale_index, view.shape)] = scale_bytes diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/indexer_k_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/indexer_k_tail.py new file mode 100644 index 000000000000..0bd627f19c12 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/indexer_k_tail.py @@ -0,0 +1,277 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Indexer K tail: k_norm, RoPE, E4M3 K and the indexer K-cache write in one Triton launch.""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._indexer_prologue import launch_indexer_k_tail +from ._indexer_reference import reference_k_rows, write_cache_rows + +#: GLM-5.2's indexer: head dim, RoPE slice (first 64 values, interleaved pairs), k_norm eps, the +#: UE8M0 scale format, and the indexer K cache's page (64 tokens of 128 FP8 bytes, then their 64 +#: fp32 scales). The pre-norm K is a column view of the fused-A output, whose rows are the +#: kv_a_proj_with_mqa (2624), indexer wk (128) and weights_proj (32) rows. +_GLM = dict( + head_dim=128, + rope_dim=64, + interleave=True, + use_ue8m0=True, + eps=1e-6, + tokens_per_block=64, + fused_a_width=2784, + k_col=2624, +) + + +def _spec(rows, first_position, **overrides) -> dict: + """``rows`` K rows of one request at positions first_position, first_position + 1, ...""" + spec = dict( + _GLM, + rows=rows, + cache_rows=rows, + first_position=first_position, + k_layout="fused_a", + pool_layout="contiguous", + pages=16, + negative_slot=False, + use_pdl=True, + launch_dependents=True, + prewait=True, + ) + spec.update(overrides) + return spec + + +class GlmIndexerKTail(OpWrapper): + """The K part of the fused indexer prologue for rows whose Q part another kernel computes. + + For every row ``i`` of ``indexer_k`` [N, 128] (the pre-norm indexer K): + + 1. ``kn = LayerNorm(bf16(indexer_k[i]); ln_weight, ln_bias, eps)``, fp32 math, rounded to bf16; + 2. RoPE on ``kn[:rope_dim]`` at ``positions[i]`` with the ``[cos | sin]`` row of ``cos_sin`` + (interleaved pairs when ``interleave``), ``fma(x, cos, (+-partner) * sin)``, rounded to bf16; + 3. ``k_scale[i] = max(amax, 1e-12) / 448``, rounded up to a power of two with ``use_ue8m0``; + ``k_fp8[i] = e4m3(k * rcp(k_scale[i]))``; + 4. for ``i < num_cache_tokens`` with both slots >= 0, the 128 FP8 bytes into + ``k_cache`` at the flat byte offset ``slot_mapping_fp8[i]`` and the 4 scale bytes at + ``slot_mapping_scale[i]``, each 4-byte chunk unravelled over ``k_cache``'s strided view + (``indexer_k_cache_scatter_op``'s rule). + + Returns ``(k_fp8, k_scale, k_cache)``, all written in place. + + With ``use_pdl`` the kernel waits on ``griddepcontrol.wait`` before reading its inputs; + ``launch_dependents`` lets the next kernel of the stream launch right after that wait (only for + the early-KV decode MQA logits, which read the new cache rows after their own wait); ``prewait`` + loads the positions, cos/sin rows, the LayerNorm parameters and the slot mappings before the + wait (only when the step's metadata preparation wrote them before the kernels this one waits + for). + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="verify / draft step 0: 8 rows from fused-A columns, PDL trigger + pre-wait", + spec=_spec(8, 529_981), + ), + Cell( + why="8 rows at the end of a 1M-token context (cos/sin rows near 2^20)", + spec=_spec(8, 1_048_560), + ), + Cell( + why="plain launch, PDL off", + spec=_spec(8, 4_093, use_pdl=False, launch_dependents=False, prewait=False), + ), + Cell( + why="PDL trigger without the pre-wait loads", + spec=_spec(8, 70_000, prewait=False), + ), + Cell( + why="padded graph rows: 8 K rows, the first 5 written to the cache", + spec=_spec(8, 130, cache_rows=5), + ), + Cell( + why="a slot of -1 leaves that row out of the cache", + spec=_spec(8, 640, negative_slot=True), + ), + Cell( + why="cache pages interleaved with other layers' pages (strided page dim)", + spec=_spec(8, 63, pool_layout="interleaved"), + ), + Cell( + why="contiguous pre-norm K (a separate wk|weights_proj GEMM)", + spec=_spec(8, 9_000, k_layout="contiguous"), + ), + ) + + note = """ + Metadata consumed: none; every input is an argument. The target passes the K columns + [2624, 2752) of the fused-A output, the step's ``position_ids`` (int32), its cos/sin cache + ``[max_pos, 64]`` fp32 (``[cos | sin]`` per row), the layer's ``k_norm`` weight and bias, the + layer's indexer K-cache view ``[pages, 64, 1, 132]`` uint8 and the step's + ``slot_mapping_fp8`` / ``slot_mapping_scale`` (int64 flat byte offsets), with + ``num_cache_tokens`` = the step's token count. + + ``k_fp8`` / ``k_scale`` are written in full (N rows) even where only the cache is read + afterwards. A slot < 0 skips the row's cache write, as the scatter op does. + + Not certified: fp32 pre-norm K, NeoX RoPE (``interleave=False``), the plain amax / 448 scale + (``use_ue8m0=False``), other head or RoPE dims, more than 8 rows. + """ + + def __call__( + self, + indexer_k: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + eps: float, + rope_dim: int, + interleave: bool, + use_ue8m0: bool, + k_fp8: torch.Tensor, + k_scale: torch.Tensor, + k_cache: torch.Tensor, + slot_mapping_fp8: torch.Tensor, + slot_mapping_scale: torch.Tensor, + num_cache_tokens: int, + use_pdl: bool, + launch_dependents: bool, + prewait: bool, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + launch_indexer_k_tail( + indexer_k, + positions, + cos_sin, + ln_weight, + ln_bias, + eps, + rope_dim, + interleave, + use_ue8m0, + k_fp8, + k_scale, + k_cache, + slot_mapping_fp8, + slot_mapping_scale, + num_cache_tokens, + use_pdl, + launch_dependents, + prewait, + ) + return k_fp8, k_scale, k_cache + + def reference( + self, + indexer_k: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + eps: float, + rope_dim: int, + interleave: bool, + use_ue8m0: bool, + k_fp8: torch.Tensor, + k_scale: torch.Tensor, + k_cache: torch.Tensor, + slot_mapping_fp8: torch.Tensor, + slot_mapping_scale: torch.Tensor, + num_cache_tokens: int, + use_pdl: bool, + launch_dependents: bool, + prewait: bool, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Fresh ``(k_fp8, k_scale, k_cache)`` from plain torch, fp32 math rounded where the kernel + rounds. PDL changes when the kernel runs, not what it computes, so the three launch flags + are unused.""" + k8, scale = reference_k_rows( + indexer_k, positions, cos_sin, ln_weight, ln_bias, eps, rope_dim, interleave, use_ue8m0 + ) + cache = k_cache.clone() + write_cache_rows(cache, k8, scale, slot_mapping_fp8, slot_mapping_scale, num_cache_tokens) + return k8, scale, cache + + def is_valid( + self, + indexer_k: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + eps: float, + rope_dim: int, + interleave: bool, + use_ue8m0: bool, + k_fp8: torch.Tensor, + k_scale: torch.Tensor, + k_cache: torch.Tensor, + slot_mapping_fp8: torch.Tensor, + slot_mapping_scale: torch.Tensor, + num_cache_tokens: int, + use_pdl: bool, + launch_dependents: bool, + prewait: bool, + ) -> None: + rows = indexer_k.shape[0] + # The kernel reads every operand through a unit inner stride and a fixed 128-wide row. + assert indexer_k.dim() == 2 and indexer_k.shape[1] == 128, ( + f"indexer_k must be [N, 128], got {tuple(indexer_k.shape)}" + ) + assert indexer_k.stride(1) == 1, "indexer_k needs a unit inner stride" + assert indexer_k.dtype in (torch.bfloat16, torch.float32), indexer_k.dtype + # The sin half starts rope_dim / 2 columns into the row: a wider or narrower table reads + # the wrong values without failing. + assert cos_sin.dtype == torch.float32 and cos_sin.dim() == 2, ( + "cos_sin must be a [max_pos, rope_dim] fp32 table" + ) + assert cos_sin.shape[1] == rope_dim and cos_sin.stride(1) == 1, ( + f"cos_sin rows must be [cos | sin] of rope_dim={rope_dim} contiguous values" + ) + assert rope_dim % 2 == 0 and 0 < rope_dim <= 128, f"rope_dim={rope_dim}" + assert positions.dim() == 1 and positions.numel() >= rows, ( + "positions must hold one position per row" + ) + if rows > 0: + assert int(positions[:rows].max()) < cos_sin.shape[0], ( + "a position past the cos/sin table reads past its end" + ) + assert ln_weight.numel() == 128 and ln_bias.numel() == 128, ( + "k_norm weight and bias hold 128 values" + ) + # The outputs are written as dense [N, 128] / [N, 1] rows. + assert k_fp8.shape == (rows, 128) and k_fp8.is_contiguous(), ( + f"k_fp8 must be a contiguous [N, 128] tensor, got {tuple(k_fp8.shape)}" + ) + assert k_fp8.dtype == torch.float8_e4m3fn, k_fp8.dtype + assert k_scale.shape == (rows, 1) and k_scale.is_contiguous(), ( + f"k_scale must be a contiguous [N, 1] tensor, got {tuple(k_scale.shape)}" + ) + assert k_scale.dtype == torch.float32, k_scale.dtype + assert k_cache.dim() == 4 and k_cache.element_size() == 1, ( + "k_cache must be the 4-D byte view of the indexer K cache" + ) + assert slot_mapping_fp8.dtype == torch.int64 and slot_mapping_scale.dtype == torch.int64, ( + "slot mappings are int64 flat byte offsets" + ) + assert 0 < num_cache_tokens <= rows, ( + f"num_cache_tokens={num_cache_tokens} must be in [1, {rows}]" + ) + assert slot_mapping_fp8.numel() >= num_cache_tokens, "too few fp8 slots" + assert slot_mapping_scale.numel() >= num_cache_tokens, "too few scale slots" + limit = k_cache.numel() + written = slot_mapping_fp8[:num_cache_tokens] + assert int(written.max()) + 128 <= limit, "an fp8 slot writes past the cache" + assert int(slot_mapping_scale[:num_cache_tokens].max()) + 4 <= limit, ( + "a scale slot writes past the cache" + ) + + +indexer_k_tail = GlmIndexerKTail() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/indexer_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/indexer_prologue.py new file mode 100644 index 000000000000..a8a9f032d90e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/indexer_prologue.py @@ -0,0 +1,272 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused indexer prologue: k_norm, RoPE, E4M3 Q and K, the scaled weights and the indexer K-cache +write in one Triton launch.""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._indexer_prologue import launch_indexer_prologue +from ._indexer_reference import reference_k_rows, reference_q_rows, write_cache_rows + +#: GLM-5.2's indexer: 32 heads of 128, the RoPE slice (first 64 values, interleaved pairs), the +#: k_norm eps, the UE8M0 scale format and softmax_scale * n_heads**-0.5. K and the weights are +#: column views of the fused-A output (kv_a_proj_with_mqa 2624 rows, wk 128, weights_proj 32). +_GLM = dict( + n_heads=32, + head_dim=128, + rope_dim=64, + interleave=True, + use_ue8m0=True, + eps=1e-6, + weight_scale=128**-0.5 * 32**-0.5, + tokens_per_block=64, + fused_a_width=2784, + k_col=2624, +) + + +def _spec(rows, first_position, **overrides) -> dict: + spec = dict( + _GLM, + rows=rows, + cache_rows=rows, + first_position=first_position, + k_layout="fused_a", + pool_layout="contiguous", + pages=16, + write_cache=True, + use_pdl=True, + launch_dependents=True, + prewait=True, + ) + spec.update(overrides) + return spec + + +class GlmIndexerPrologue(OpWrapper): + """The whole indexer prologue of a step whose Q comes from a separate wq_b GEMM. + + For every row ``i`` (the K part is ``indexer_k_tail``'s, including the cache write when + ``k_cache`` is given and ``num_cache_tokens`` > 0): + + - Q: RoPE on ``q[i, h, :rope_dim]`` at ``positions[i]``, rounded to bf16; per head + ``q_scale[i, h] = max(amax, 1e-12) / 448`` (rounded up to a power of two with + ``use_ue8m0``) and ``q_fp8[i, h] = e4m3(q * rcp(q_scale))``; + - weights: ``(weights[i, h] * q_scale[i, h]) * weight_scale`` in fp32. + + Returns ``(q_fp8 [N, H, D], k_fp8 [N, D], k_scale [N, 1], weights [N, H] fp32, + q_scale [N, H, 1])``, freshly allocated, the contract of the indexer's Q and K inputs to the + MQA logits. PDL flags as in ``indexer_k_tail``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode, 8 rows: wq_b output Q, fused-A K/weights, cache write, PDL + pre-wait", + spec=_spec(8, 529_981), + ), + Cell( + why="decode, 8 rows, plain launch", + spec=_spec(8, 1_048_560, use_pdl=False, launch_dependents=False, prewait=False), + ), + Cell( + why="no cache write (the caller scatters the K rows itself)", + spec=_spec(8, 300, write_cache=False, launch_dependents=False, prewait=False), + ), + Cell( + why="padded graph rows: 8 rows, the first 3 written to the cache", + spec=_spec(8, 60, cache_rows=3), + ), + Cell( + why="context chunk: 256 rows of one request over 4+ pages, contiguous K and weights", + spec=_spec( + 256, + 1_000, + k_layout="contiguous", + use_pdl=False, + launch_dependents=False, + prewait=False, + ), + ), + ) + + note = """ + Metadata consumed: none; every input is an argument (see ``indexer_k_tail`` for the K-side + operands). ``q`` is the indexer wq_b output viewed as [N, 32, 128] bf16 and ``weights`` the + weights_proj columns [2752, 2784) of the fused-A output (bf16). + + At the c1 decode cells the target takes the Q part from the wq_b epilogue and runs only + ``indexer_k_tail``; this entry serves the forwards that epilogue does not. + + Not certified: fp32 K or weights, NeoX RoPE, the plain amax / 448 scale, other head counts or + dims. + """ + + def __call__( + self, + indexer_k: torch.Tensor, + weights: torch.Tensor, + q: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + eps: float, + weight_scale: float, + rope_dim: int, + interleave: bool, + use_ue8m0: bool, + k_cache: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + num_cache_tokens: int, + use_pdl: bool, + launch_dependents: bool, + prewait: bool, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return launch_indexer_prologue( + indexer_k, + weights, + q, + positions, + cos_sin, + ln_weight, + ln_bias, + eps, + weight_scale, + rope_dim, + interleave, + use_ue8m0, + k_cache, + slot_mapping_fp8, + slot_mapping_scale, + num_cache_tokens, + use_pdl, + launch_dependents, + prewait, + ) + + def reference( + self, + indexer_k: torch.Tensor, + weights: torch.Tensor, + q: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + eps: float, + weight_scale: float, + rope_dim: int, + interleave: bool, + use_ue8m0: bool, + k_cache: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + num_cache_tokens: int, + use_pdl: bool, + launch_dependents: bool, + prewait: bool, + ) -> Tuple[torch.Tensor, ...]: + """``(q_fp8, k_fp8, k_scale, weights, q_scale, k_cache)``: the five outputs, then the cache + after the writes (a fresh tensor; None without ``k_cache``). The launch flags change when + the kernel runs, not what it computes.""" + q8, q_scale, w = reference_q_rows( + q, weights, positions, cos_sin, weight_scale, rope_dim, interleave, use_ue8m0 + ) + k8, k_scale = reference_k_rows( + indexer_k, positions, cos_sin, ln_weight, ln_bias, eps, rope_dim, interleave, use_ue8m0 + ) + cache = None + if k_cache is not None: + cache = k_cache.clone() + if num_cache_tokens > 0: + write_cache_rows( + cache, k8, k_scale, slot_mapping_fp8, slot_mapping_scale, num_cache_tokens + ) + return q8, k8, k_scale, w, q_scale, cache + + def is_valid( + self, + indexer_k: torch.Tensor, + weights: torch.Tensor, + q: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + eps: float, + weight_scale: float, + rope_dim: int, + interleave: bool, + use_ue8m0: bool, + k_cache: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + num_cache_tokens: int, + use_pdl: bool, + launch_dependents: bool, + prewait: bool, + ) -> None: + rows, n_heads, head_dim = q.shape + # Every operand is read through a unit inner stride and fixed 128-wide rows. + assert head_dim == 128 and indexer_k.shape == (rows, 128), ( + f"q must be [N, H, 128] and indexer_k [N, 128], got {tuple(q.shape)} and " + f"{tuple(indexer_k.shape)}" + ) + assert q.dtype == torch.bfloat16 and q.stride(2) == 1, "q must be bf16, unit inner stride" + assert indexer_k.stride(1) == 1 and weights.stride(1) == 1, ( + "indexer_k and weights need a unit inner stride" + ) + assert weights.shape == (rows, n_heads), ( + f"weights must be [N, H] = [{rows}, {n_heads}], got {tuple(weights.shape)}" + ) + assert indexer_k.dtype in (torch.bfloat16, torch.float32), indexer_k.dtype + assert weights.dtype in (torch.bfloat16, torch.float32), weights.dtype + assert cos_sin.dtype == torch.float32 and cos_sin.dim() == 2, ( + "cos_sin must be a [max_pos, rope_dim] fp32 table" + ) + # The sin half starts rope_dim / 2 columns into the row. + assert cos_sin.shape[1] == rope_dim and cos_sin.stride(1) == 1, ( + f"cos_sin rows must be [cos | sin] of rope_dim={rope_dim} contiguous values" + ) + assert rope_dim % 2 == 0 and 0 < rope_dim <= 128, f"rope_dim={rope_dim}" + assert positions.dim() == 1 and positions.numel() >= rows, ( + "positions must hold one position per row" + ) + if rows > 0: + assert int(positions[:rows].max()) < cos_sin.shape[0], ( + "a position past the cos/sin table reads past its end" + ) + assert ln_weight.numel() == 128 and ln_bias.numel() == 128, ( + "k_norm weight and bias hold 128 values" + ) + if k_cache is not None and num_cache_tokens > 0: + assert k_cache.dim() == 4 and k_cache.element_size() == 1, ( + "k_cache must be the 4-D byte view of the indexer K cache" + ) + assert slot_mapping_fp8 is not None and slot_mapping_scale is not None, ( + "a cache write needs both slot mappings" + ) + assert slot_mapping_fp8.dtype == torch.int64, "slot mappings are int64" + assert slot_mapping_scale.dtype == torch.int64, "slot mappings are int64" + assert num_cache_tokens <= rows, f"num_cache_tokens={num_cache_tokens} > {rows} rows" + assert slot_mapping_fp8.numel() >= num_cache_tokens, "too few fp8 slots" + assert slot_mapping_scale.numel() >= num_cache_tokens, "too few scale slots" + limit = k_cache.numel() + assert int(slot_mapping_fp8[:num_cache_tokens].max()) + 128 <= limit, ( + "an fp8 slot writes past the cache" + ) + assert int(slot_mapping_scale[:num_cache_tokens].max()) + 4 <= limit, ( + "a scale slot writes past the cache" + ) + + +indexer_prologue = GlmIndexerPrologue() diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_indexer_inputs.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_indexer_inputs.py new file mode 100644 index 000000000000..ce7f6d28e8e4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_indexer_inputs.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Inputs the GLM indexer entries' tests share: the indexer K cache, its slot mappings and block +table, the RoPE table, and the comparison of E4M3 rows quantized with per-row scales. + +The cache layout is the engine's: an INDEX_KEY page holds 64 tokens' 128 FP8 bytes, then their 64 +fp32 scales, viewed as [pages, 64, 1, 132] uint8; a token's bytes are found through the slot +mappings, flat byte offsets computed exactly as the DSA metadata does. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +PAGE = 64 +HEAD_DIM = 128 +TOKEN_BYTES = HEAD_DIM + 4 +SENTINEL = 0xA5 + + +def cos_sin_table(rows: int, rope_dim: int = 64, theta: float = 8.0e6) -> torch.Tensor: + """[rows, rope_dim] fp32, ``[cos | sin]`` of the default RoPE at every position.""" + exponents = torch.arange(0, rope_dim, 2, dtype=torch.float64, device="cuda") / rope_dim + inv_freq = 1.0 / (theta**exponents) + angles = torch.arange(rows, dtype=torch.float64, device="cuda")[:, None] * inv_freq[None, :] + return torch.cat([angles.cos(), angles.sin()], dim=1).float() + + +def k_cache_pool(pages: int, layout: str) -> Tuple[torch.Tensor, torch.Tensor]: + """(this layer's [pages, 64, 1, 132] uint8 view, the whole allocation), filled with a + sentinel. ``interleaved`` puts two other layers' pages between this layer's pages (a strided + page dimension); ``contiguous`` is one layer's own pool.""" + if layout == "interleaved": + whole = torch.full((pages, 3, PAGE * TOKEN_BYTES), SENTINEL, dtype=torch.uint8).cuda() + view = whole[:, 1].view(pages, PAGE, 1, TOKEN_BYTES) + assert not view.is_contiguous() + return view, whole + whole = torch.full((pages, PAGE, 1, TOKEN_BYTES), SENTINEL, dtype=torch.uint8).cuda() + return whole, whole + + +def block_row(positions: torch.Tensor, pool_pages: int, gen: torch.Generator) -> torch.Tensor: + """One request's INDEX_KEY block-table row up to the page of its last position: the pages the + positions touch map to distinct random physical pages (so no two rows share a slot), the + others to page 0.""" + pages = positions.long() // PAGE + touched = torch.unique(pages) + assert touched.numel() <= pool_pages, "the pool is smaller than the pages the rows touch" + row = torch.zeros(int(pages.max()) + 1, dtype=torch.int32) + row[touched] = torch.randperm(pool_pages, generator=gen)[: touched.numel()].to(torch.int32) + return row + + +def slot_mappings( + positions: torch.Tensor, row: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + """(fp8 slots, scale slots) int64 flat byte offsets of each position's K row, the metadata's + ``_compute_slot_mappings`` for one request.""" + pos = positions.long().cpu() + block_ids = row.long()[pos // PAGE] + in_page = pos % PAGE + page_stride = PAGE * TOKEN_BYTES + fp8 = block_ids * page_stride + in_page * HEAD_DIM + scale = block_ids * page_stride + PAGE * HEAD_DIM + in_page * 4 + return fp8.cuda(), scale.cuda() + + +def e4m3_step(units: torch.Tensor) -> torch.Tensor: + """Spacing of E4M3 values at magnitude ``units`` (in units of the scale).""" + exponent = torch.floor(torch.log2(units.clamp_min(2.0**-6))) + return torch.pow(2.0, exponent - 3) + + +def assert_quantized_close( + q8: torch.Tensor, + scale: torch.Tensor, + q8_ref: torch.Tensor, + scale_ref: torch.Tensor, + max_mismatch: float = 0.01, +) -> None: + """E4M3 rows with per-row scales against a reference that rounds at the same points. + + The kernel and the reference compute the same fp32 expressions, but the LayerNorm reduction + order and the MUFU rsqrt can move an fp32 intermediate by an ulp, and a bf16 rounding of that + intermediate then flips with probability of about 2^-16 per value. A flipped bf16 value moves + its E4M3 code by at most one step (and the row scale, a power of two, only when the row's + maximum crosses one). So every dequantized value is within one E4M3 step of the reference, at + the coarser of the two scales, and almost all codes match exactly. + """ + assert q8.shape == q8_ref.shape and scale.shape == scale_ref.shape + deq = q8.float() * scale + deq_ref = q8_ref.float() * scale_ref + step = torch.maximum( + scale * e4m3_step(q8.float().abs()), scale_ref * e4m3_step(q8_ref.float().abs()) + ) + excess = ((deq - deq_ref).abs() - step).clamp_min(0) + assert excess.max().item() == 0, ( + f"a value is more than one E4M3 step from the reference (by {excess.max().item():.3e})" + ) + mismatch = (q8.view(torch.uint8) != q8_ref.view(torch.uint8)).float().mean().item() + assert mismatch <= max_mismatch, f"{mismatch:.2%} of the E4M3 codes differ" diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_indexer_k_tail.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_indexer_k_tail.py new file mode 100644 index 000000000000..534f616dc582 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_indexer_k_tail.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM indexer_k_tail entry. + +The K cache is real: an INDEX_KEY pool of 64-token pages ([pages, 64, 1, 132] uint8) filled with a +sentinel, a random block-table row for the request and the metadata's slot mappings. The cache is +checked in two ways: bitwise against the op's own K outputs scattered by the reference's byte +placement (so every byte the op wrote is where it belongs, and nothing else moved), and through +those outputs against the reference arithmetic. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa import ( # noqa: E501 + _indexer_reference, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.indexer_k_tail import ( # noqa: E501 + indexer_k_tail as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../..", "."] +import _indexer_inputs as inputs # noqa: E402 +from _validating import certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +assert torch.cuda.is_available(), "indexer_k_tail requires a CUDA device" + + +def _build(spec, seed: int) -> dict: + """The op's arguments for one cell. The whole pool allocation rides along as ``_pool``.""" + gen = torch.Generator().manual_seed(seed) + rows, d = spec["rows"], spec["head_dim"] + positions = torch.arange(spec["first_position"], spec["first_position"] + rows) + if spec["k_layout"] == "fused_a": + fused = torch.randn(rows, spec["fused_a_width"], generator=gen).to(torch.bfloat16).cuda() + indexer_k = fused[:, spec["k_col"] : spec["k_col"] + d] + assert not indexer_k.is_contiguous() + else: + indexer_k = torch.randn(rows, d, generator=gen).to(torch.bfloat16).cuda() + ln_weight = (1.0 + 0.1 * torch.randn(d, generator=gen)).to(torch.bfloat16).cuda() + ln_bias = (0.1 * torch.randn(d, generator=gen)).to(torch.bfloat16).cuda() + cache, whole = inputs.k_cache_pool(spec["pages"], spec["pool_layout"]) + row = inputs.block_row(positions, spec["pages"], gen) + slot_fp8, slot_scale = inputs.slot_mappings(positions, row) + if spec["negative_slot"]: + slot_fp8[2] = -1 + return dict( + _pool=whole, + indexer_k=indexer_k, + positions=positions.to(torch.int32).cuda(), + cos_sin=inputs.cos_sin_table(int(positions[-1]) + 1, spec["rope_dim"]), + ln_weight=ln_weight, + ln_bias=ln_bias, + eps=spec["eps"], + rope_dim=spec["rope_dim"], + interleave=spec["interleave"], + use_ue8m0=spec["use_ue8m0"], + k_fp8=torch.empty(rows, d, dtype=torch.float8_e4m3fn, device="cuda"), + k_scale=torch.empty(rows, 1, dtype=torch.float32, device="cuda"), + k_cache=cache, + slot_mapping_fp8=slot_fp8, + slot_mapping_scale=slot_scale, + num_cache_tokens=spec["cache_rows"], + use_pdl=spec["use_pdl"], + launch_dependents=spec["launch_dependents"], + prewait=spec["prewait"], + ) + + +def _check(args: dict, out, ref, cache_before: torch.Tensor, whole_before=None, whole=None): + k8, k_scale, cache = out + k8_ref, k_scale_ref, _ = ref + # The arithmetic: within one E4M3 step of the reference, almost every code exact. + inputs.assert_quantized_close(k8, k_scale, k8_ref, k_scale_ref) + # The placement: the op's own rows, scattered by the reference rule into the cache as it was, + # are the cache byte for byte. + expected = cache_before.clone() + _indexer_reference.write_cache_rows( + expected, + k8, + k_scale, + args["slot_mapping_fp8"], + args["slot_mapping_scale"], + args["num_cache_tokens"], + ) + assert torch.equal(cache.view(torch.uint8), expected.view(torch.uint8)), "cache bytes differ" + if whole is not None: + # The other layers' pages of an interleaved pool are untouched. + assert torch.equal(whole[:, 0], whole_before[:, 0]) + assert torch.equal(whole[:, 2], whole_before[:, 2]) + + +def _op_args(args: dict) -> dict: + return {k: v for k, v in args.items() if not k.startswith("_")} + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + whole = args.pop("_pool") if cell.spec["pool_layout"] == "interleaved" else None + args = _op_args(args) + cache_before = args["k_cache"].clone() + whole_before = whole.clone() if whole is not None else None + ref = op.reference(**args) + out = op(**args) + torch.cuda.synchronize() + assert out[0] is args["k_fp8"] and out[1] is args["k_scale"] and out[2] is args["k_cache"] + _check(args, out, ref, cache_before, whole_before, whole) + + +def test_rows_past_num_cache_tokens_stay_out_of_the_cache(): + """Rows at and past ``num_cache_tokens`` get K outputs and no cache bytes.""" + spec = dict(op.CELLS[0].spec, cache_rows=2) + args = _op_args(_build(spec, seed=11)) + before = args["k_cache"].clone() + op(**args) + torch.cuda.synchronize() + changed = (args["k_cache"].view(torch.uint8) != before.view(torch.uint8)).sum().item() + # Two rows: 128 FP8 bytes and 4 scale bytes each, at most (a value may equal the sentinel). + assert 0 < changed <= 2 * (128 + 4) + + +def test_cuda_graph_replay_follows_new_inputs(): + """The decode path is captured: a replay reads the K rows, positions and slots in place.""" + spec = op.CELLS[0].spec + args = _op_args(_build(spec, seed=21)) + op(**args) # compile outside capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(**args) + # Positions inside the captured cos/sin table, rows of another request map. + fresh = _build(dict(spec, first_position=spec["first_position"] - 8), seed=22) + args["indexer_k"].copy_(fresh["indexer_k"]) + args["positions"].copy_(fresh["positions"]) + args["slot_mapping_fp8"].copy_(fresh["slot_mapping_fp8"]) + args["slot_mapping_scale"].copy_(fresh["slot_mapping_scale"]) + cache_before = args["k_cache"].clone() + ref = op.reference(**args) + graph.replay() + torch.cuda.synchronize() + _check(args, (args["k_fp8"], args["k_scale"], args["k_cache"]), ref, cache_before) + + +def test_guard_refuses_a_wider_cos_sin_table(): + args = _op_args(_build(op.CELLS[0].spec, seed=32)) + wide = torch.cat([args["cos_sin"], args["cos_sin"]], dim=1) + with pytest.raises(AssertionError, match="cos_sin rows"): + with validating(op): + op(**dict(args, cos_sin=wide)) + + +def test_guard_refuses_a_position_past_the_table(): + args = _op_args(_build(op.CELLS[0].spec, seed=33)) + positions = args["positions"].clone() + positions[-1] = args["cos_sin"].shape[0] + with pytest.raises(AssertionError, match="past the cos/sin table"): + with validating(op): + op(**dict(args, positions=positions)) + + +def test_guard_refuses_a_strided_k_output(): + args = _op_args(_build(op.CELLS[0].spec, seed=34)) + wide = torch.empty(args["k_fp8"].shape[0], 256, dtype=torch.float8_e4m3fn, device="cuda") + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op(**dict(args, k_fp8=wide[:, :128])) + + +def test_guard_refuses_a_slot_past_the_cache(): + args = _op_args(_build(op.CELLS[0].spec, seed=35)) + slots = args["slot_mapping_fp8"].clone() + slots[0] = args["k_cache"].numel() - 64 + with pytest.raises(AssertionError, match="writes past the cache"): + with validating(op): + op(**dict(args, slot_mapping_fp8=slots)) + + +def test_guard_refuses_more_cache_rows_than_k_rows(): + args = _op_args(_build(op.CELLS[0].spec, seed=36)) + with pytest.raises(AssertionError, match="num_cache_tokens"): + with validating(op): + op(**dict(args, num_cache_tokens=args["indexer_k"].shape[0] + 1)) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_indexer_prologue.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_indexer_prologue.py new file mode 100644 index 000000000000..e702c25a0f32 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_indexer_prologue.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM indexer_prologue entry. + +Q is the wq_b output ([N, 32 * 128] bf16); K and the weights are column views of a fused-A output +row block, as the target hands them over. The cache is the engine's INDEX_KEY layout (see +``_indexer_inputs``); its bytes are checked against the op's own K rows scattered by the reference +rule, the outputs against the reference arithmetic. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa import ( # noqa: E501 + _indexer_reference, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.indexer_prologue import ( # noqa: E501 + indexer_prologue as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../..", "."] +import _indexer_inputs as inputs # noqa: E402 +from _validating import certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +assert torch.cuda.is_available(), "indexer_prologue requires a CUDA device" + + +def _build(spec, seed: int) -> dict: + gen = torch.Generator().manual_seed(seed) + rows, heads, d = spec["rows"], spec["n_heads"], spec["head_dim"] + positions = torch.arange(spec["first_position"], spec["first_position"] + rows) + q = torch.randn(rows, heads * d, generator=gen).to(torch.bfloat16).cuda() + if spec["k_layout"] == "fused_a": + fused = torch.randn(rows, spec["fused_a_width"], generator=gen).to(torch.bfloat16).cuda() + indexer_k = fused[:, spec["k_col"] : spec["k_col"] + d] + weights = fused[:, spec["k_col"] + d : spec["k_col"] + d + heads] + else: + indexer_k = torch.randn(rows, d, generator=gen).to(torch.bfloat16).cuda() + weights = torch.randn(rows, heads, generator=gen).to(torch.bfloat16).cuda() + cache = slot_fp8 = slot_scale = None + if spec["write_cache"]: + cache, _ = inputs.k_cache_pool(spec["pages"], spec["pool_layout"]) + row = inputs.block_row(positions, spec["pages"], gen) + slot_fp8, slot_scale = inputs.slot_mappings(positions, row) + return dict( + indexer_k=indexer_k, + weights=weights, + q=q.view(rows, heads, d), + positions=positions.to(torch.int32).cuda(), + cos_sin=inputs.cos_sin_table(int(positions[-1]) + 1, spec["rope_dim"]), + ln_weight=(1.0 + 0.1 * torch.randn(d, generator=gen)).to(torch.bfloat16).cuda(), + ln_bias=(0.1 * torch.randn(d, generator=gen)).to(torch.bfloat16).cuda(), + eps=spec["eps"], + weight_scale=spec["weight_scale"], + rope_dim=spec["rope_dim"], + interleave=spec["interleave"], + use_ue8m0=spec["use_ue8m0"], + k_cache=cache, + slot_mapping_fp8=slot_fp8, + slot_mapping_scale=slot_scale, + num_cache_tokens=spec["cache_rows"] if spec["write_cache"] else 0, + use_pdl=spec["use_pdl"], + launch_dependents=spec["launch_dependents"], + prewait=spec["prewait"], + ) + + +def _check(args, out, ref, cache_before): + q8, k8, k_scale, w, q_scale = out + q8_ref, k8_ref, k_scale_ref, w_ref, q_scale_ref, _ = ref + rows, heads, d = args["q"].shape + assert q8.shape == (rows, heads, d) and q_scale.shape == (rows, heads, 1) + assert w.shape == (rows, heads) and w.dtype == torch.float32 + inputs.assert_quantized_close(q8, q_scale, q8_ref, q_scale_ref) + inputs.assert_quantized_close(k8, k_scale, k8_ref, k_scale_ref) + # weights = (w * q_scale) * weight_scale, two fp32 roundings: dividing the Q scale back out + # leaves the input weight times the constant within two ulps whichever scale either side + # picked. + torch.testing.assert_close( + w / q_scale[..., 0], + w_ref / q_scale_ref[..., 0], + rtol=2 * 2.0**-23, + atol=0.0, + ) + if args["k_cache"] is not None: + expected = cache_before.clone() + _indexer_reference.write_cache_rows( + expected, + k8, + k_scale, + args["slot_mapping_fp8"], + args["slot_mapping_scale"], + args["num_cache_tokens"], + ) + assert torch.equal(args["k_cache"], expected), "cache bytes differ" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + cache_before = args["k_cache"].clone() if args["k_cache"] is not None else None + ref = op.reference(**args) + out = op(**args) + torch.cuda.synchronize() + _check(args, out, ref, cache_before) + + +def test_q_part_matches_the_k_tail_launch_of_the_same_kernel(): + """The K part does not depend on whether the Q part is compiled in: the whole prologue and the + K tail write the same K bytes for the same inputs.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.indexer_k_tail import ( # noqa: E501 + indexer_k_tail, + ) + + args = _build(op.CELLS[0].spec, seed=41) + _, k8, k_scale, _, _ = op(**dict(args, k_cache=None, num_cache_tokens=0)) + tail_k8 = torch.empty_like(k8) + tail_scale = torch.empty_like(k_scale) + indexer_k_tail( + indexer_k=args["indexer_k"], + positions=args["positions"], + cos_sin=args["cos_sin"], + ln_weight=args["ln_weight"], + ln_bias=args["ln_bias"], + eps=args["eps"], + rope_dim=args["rope_dim"], + interleave=args["interleave"], + use_ue8m0=args["use_ue8m0"], + k_fp8=tail_k8, + k_scale=tail_scale, + k_cache=args["k_cache"], + slot_mapping_fp8=args["slot_mapping_fp8"], + slot_mapping_scale=args["slot_mapping_scale"], + num_cache_tokens=args["num_cache_tokens"], + use_pdl=args["use_pdl"], + launch_dependents=args["launch_dependents"], + prewait=args["prewait"], + ) + torch.cuda.synchronize() + assert torch.equal(k8.view(torch.uint8), tail_k8.view(torch.uint8)) + assert torch.equal(k_scale, tail_scale) + + +def test_guard_refuses_weights_of_another_head_count(): + args = _build(op.CELLS[0].spec, seed=42) + with pytest.raises(AssertionError, match="weights must be"): + with validating(op): + op(**dict(args, weights=args["weights"][:, :16])) + + +def test_guard_refuses_a_cache_write_without_slots(): + args = _build(op.CELLS[0].spec, seed=43) + with pytest.raises(AssertionError, match="both slot mappings"): + with validating(op): + op(**dict(args, slot_mapping_scale=None)) + + +def test_guard_refuses_a_non_unit_inner_stride_q(): + args = _build(op.CELLS[0].spec, seed=44) + rows, heads, d = args["q"].shape + strided = torch.empty(rows, heads, 2 * d, dtype=torch.bfloat16, device="cuda")[..., ::2] + with pytest.raises(AssertionError, match="unit inner stride"): + with validating(op): + op(**dict(args, q=strided)) From 4cc8374f48d5b6c68fc0290ae3d5f90a4282617d Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:19:50 -0700 Subject: [PATCH 32/53] Add the GLM FP8 paged MQA logits with the full-query and early-KV kernels A target-local copy of the CuTe DSL FP8 paged MQA logits kernel scores a request's whole verify window as one MMA atom of next_n * num_heads columns, with the weight register cache capped so that the shared-memory tail serves the remaining heads in the same FMA order. Its early-KV variant splits each group's accumulator into two N halves and is launched with PDL: it loads the KV blocks older than each request's newest next_n keys before griddepcontrol.wait, Q, the weights and the rest after it. All variants compute every element in the same order. The runner takes the caller's padded logits buffer, so captured graphs keep its address. The tp16 target gets two entries: fp8_paged_mqa_logits and paged_mqa_logits_schedule, the per-SM partition of the requests' 256-key splits that DeepGEMM's paged MQA-logits metadata computes. Both are certified on sm_107: the logits bit-exactly on integer data and within a derived band on real values, the early-KV kernel bitwise equal to the default one. Signed-off-by: Dongfeng Yu --- .../ops/dsa/_fp8_paged_mqa_logits.py | 194 ++ .../ops/dsa/_fp8_paged_mqa_logits_kernel.py | 2088 +++++++++++++++++ .../ops/dsa/fp8_paged_mqa_logits.py | 250 ++ .../ops/dsa/paged_mqa_logits_schedule.py | 110 + .../test_modeling_v2_fp8_paged_mqa_logits.py | 202 ++ ...t_modeling_v2_paged_mqa_logits_schedule.py | 76 + 6 files changed, 2920 insertions(+) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/fp8_paged_mqa_logits.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/paged_mqa_logits_schedule.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_fp8_paged_mqa_logits.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_paged_mqa_logits_schedule.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits.py new file mode 100644 index 000000000000..84b4f42fe54b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits.py @@ -0,0 +1,194 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Compile cache and launch of the indexer's FP8 paged MQA logits kernel (FP32 epilogue). + +The kernel (``_fp8_paged_mqa_logits_kernel.FP8MQALogitsKernel``) scores every query row of a +request against the request's paged indexer keys: ``logits[b * next_n + t, j] = scale[j] * +sum_h relu(q[b, t, h] . k[j]) * weights[b * next_n + t, h]`` for the keys ``j`` of the request's +context. Two policies decide the compiled variant: + +- full query: a request's ``next_n`` query rows stay one atom of ``next_n * num_heads`` MMA + columns (``full_query_weight_cache_limit``); the weight register cache is capped so that the + shared-memory tail handles the remaining heads without changing the FMA order; +- early KV (``early_kv``, shapes of ``early_kv_supported``): the N-split variant launched with + PDL, which loads the older KV blocks before ``griddepcontrol.wait``. + +All variants compute the same per-element arithmetic in the same order. +""" + +from __future__ import annotations + +import functools +from typing import Dict, Optional, Tuple + +import cutlass +import cutlass.cute as cute +import torch + +from ._fp8_paged_mqa_logits_kernel import FP8MQALogitsKernel + +#: Keys per MMA tile; a schedule split is one tile per math warp group. +COMPUTE_BLOCK_KV = 128 +#: Keys per schedule split (DeepGEMM's SPLIT_KV of the SM100 paged metadata). +SPLIT_KV = 2 * COMPUTE_BLOCK_KV + +_KERNELS: Dict[Tuple, object] = {} + + +@functools.lru_cache(maxsize=None) +def num_sms(device_index: int) -> int: + """SMs of the device, the number of persistent CTAs the kernel launches.""" + return torch.cuda.get_device_properties(device_index).multi_processor_count + + +def full_query_weight_cache_limit( + next_n: int, num_heads: int, head_dim: int, phys_block_kv: int, q_dtype: torch.dtype +) -> Optional[int]: + """The weight register cache of the full-query kernel, or None for the default kernel. + + The two MMA groups share 512 TMEM columns. Each needs ``next_n * num_heads`` columns, with an + eight-column MMA granularity. D = 128 and three Q/KV stages fit the shared-memory budget + throughout this range. Physical pages must divide the 128-token compute tile into at most four + pieces. At most 96 FP32 weights per thread (16 for each query row) stay in registers; the + shared-memory tail handles the remaining heads without changing the FMA order. Up to four + query rows keep the default kernel. + """ + columns = next_n * num_heads + if ( + next_n <= 4 + or head_dim != 128 + or not 8 <= columns <= 256 + or columns % 8 + or phys_block_kv not in (32, 64, 128) + or num_heads % 4 + or q_dtype != torch.float8_e4m3fn + ): + return None + limit = min(16, num_heads, (96 // next_n // 4) * 4) + return limit if limit >= 4 else None + + +def early_kv_supported(next_n: int, num_heads: int) -> bool: + """Shapes of the early-KV kernel (PDL launch, N split): each group's accumulator splits into + two N halves of next_n / 2 query rows, and a half must be a whole number of 16-column UMMA + atoms.""" + return next_n % 2 == 0 and (next_n * num_heads // 2) % 16 == 0 + + +def _compile( + phys_block_kv: int, + num_heads: int, + head_dim: int, + next_n: int, + sms: int, + weight_cache_limit: Optional[int], + early_kv: bool, +): + """Compile one variant against fake tensors with symbolic runtime extents (TVM FFI).""" + n = next_n * num_heads + block_bytes = phys_block_kv * (head_dim + 4) + sym_num_phys_blocks = cute.sym_int() + sym_b = cute.sym_int() + max_ctx = cute.sym_int() + max_blocks_per_seq = cute.sym_int() + num_ctas = cute.sym_int() + + # The indexer K-cache pool view can be strided in dim 0 (pages of other layers in between): + # the per-page stride is read at runtime; the bytes of one page are contiguous. + kv_fake = cute.runtime.make_fake_tensor( + cutlass.Uint8, (sym_num_phys_blocks, block_bytes), stride=(cute.sym_int64(), 1) + ) + q_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Uint8, (n, head_dim, sym_b), stride_order=(1, 0, 2) + ) + w_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (n, sym_b), stride_order=(0, 1) + ) + logits_fake = cute.runtime.make_fake_tensor( + cutlass.Float32, (cute.sym_int(), max_ctx), stride=(cute.sym_int64(), 1) + ) + bt_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_b, max_blocks_per_seq), stride_order=(1, 0) + ) + cl_fake = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (sym_b,), stride_order=(0,)) + sm_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (num_ctas, 2), stride_order=(1, 0) + ) + fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + + kernel = FP8MQALogitsKernel( + block_kv=COMPUTE_BLOCK_KV, + phys_block_kv=phys_block_kv, + num_heads=num_heads, + head_dim=head_dim, + next_n=next_n, + num_sms=sms, + num_epi_subtiles=1, + epi_dtype=cutlass.Float32, + acc_dtype=cutlass.Float32, + output_dtype=cutlass.Float32, + weight_cache_limit=weight_cache_limit, + split_n=early_kv, + early_kv=early_kv, + ) + return cute.compile( + kernel, + kv_fake, + q_fake, + w_fake, + logits_fake, + bt_fake, + cl_fake, + sm_fake, + cutlass.Int32(1), + cutlass.Int32(1), + fake_stream, + options="--enable-tvm-ffi", + ) + + +def run_fp8_paged_mqa_logits( + q: torch.Tensor, + kv_cache: torch.Tensor, + weights: torch.Tensor, + context_lens: torch.Tensor, + block_table: torch.Tensor, + schedule_meta: torch.Tensor, + logits: torch.Tensor, + max_context_len: int, + early_kv: bool, +) -> torch.Tensor: + """Launch the kernel; returns ``logits[:, :max_context_len]``. + + Args: + q: [B, next_n, H, D] float8_e4m3fn. + kv_cache: [num_pages, page_size, 1, D + 4] uint8 indexer K-cache view. + weights: [B * next_n, H] float32. + context_lens: [B] int32 keys per request. + block_table: [B, max_pages] int32, contiguous. + schedule_meta: [num_sms + 1, 2] int32 (the SM100 paged MQA-logits schedule). + logits: [B * next_n, >= round_up(max_context_len, SPLIT_KV)] float32 output, unit + column stride; the kernel writes whole splits. + max_context_len: columns of the returned view. + early_kv: run the early-KV variant where its shape allows (``early_kv_supported``). + """ + b, next_n, h, d = q.shape + n = next_n * h + phys_block_kv = kv_cache.shape[1] + num_phys_blocks = kv_cache.shape[0] + sms = num_sms(q.device.index if q.device.index is not None else torch.cuda.current_device()) + weight_cache_limit = full_query_weight_cache_limit(next_n, h, d, phys_block_kv, q.dtype) + early_kv = early_kv and early_kv_supported(next_n, h) + key = (phys_block_kv, h, d, next_n, sms, weight_cache_limit, early_kv) + compiled = _KERNELS.get(key) + if compiled is None: + compiled = _KERNELS[key] = _compile( + phys_block_kv, h, d, next_n, sms, weight_cache_limit, early_kv + ) + # Q [B, next_n, H, D] -> [N, D, B]; weights [B * next_n, H] -> [N, B]; FP8 Q as uint8 bytes. + q_3d = q.reshape(b, n, d).permute(1, 2, 0).view(torch.uint8) + w_2d = weights.reshape(b, n).t() + kv_flat = kv_cache.reshape(num_phys_blocks, -1) + out = logits[:, :max_context_len] + compiled(kv_flat, q_3d, w_2d, out, block_table, context_lens, schedule_meta, num_phys_blocks, b) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits_kernel.py new file mode 100644 index 000000000000..ca300c385d6e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits_kernel.py @@ -0,0 +1,2088 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +""" +CuTe DSL FP8 paged MQA logits kernel for Blackwell (SM100 family), with the N split and the early KV +loads of the GLM indexer decode. + +Architecture: + - 384 threads: 256 math (2 WGs) + 128 specialized (2 TMA + 2 UMMA) + - 1 TMA per KV block [128, 128], UMMA iterates 4x K=32 + - 2 warp groups process 2 KV blocks per iteration (kNumMathWarpGroups=2) + - Q reloaded via TMA pipeline when q_idx (batch) changes + - Persistent kernel: CTAs iterate through assigned (q_idx, kv_idx) pairs + - Weights cached in registers: preloaded once per q_idx change (not per KV block) + - KV Scales loaded via TMA to SMEM (separate pipeline per group, Math consumes) + +Merged KV+Scale pipeline: + - KV data and scales share a single TMA barrier per group + - TMA loads both KV and Scale under one barrier (combined tx_count) + - UMMA waits on merged barrier (for KV GEMM), does NOT release + - Math waits on merged barrier (for scale read), Math releases + +Fused KV layout: + - KV data and scales stored contiguously per physical block: + [num_phys_blocks, block_kv * (head_dim + 4)] bytes + - Per block: [KV_all_tokens (block_kv * head_dim bytes)] [Scales (block_kv * 4 bytes)] + - KV and Scale views are derived inside __call__ using CuTE pointer arithmetic + +Scheduler: + - schedule_meta[sm_idx] = (start_q_idx, start_kv_idx / kNumMathWarpGroups) + - schedule_meta[sm_idx+1] = end boundary for this CTA + - fetch_next_task pattern: each warp role independently advances (q_idx, kv_idx) + - kv_idx in units of KV blocks, advances by kNumMathWarpGroups=2 per step + +N split (split_n): + - A group's accumulator of one KV block (next_n * num_heads columns) is computed in two + N halves (query rows [0, next_n/2) and [next_n/2, next_n)), each its own TMEM stage + with its own UMMA commit + - Q is loaded as one TMA box per half; the UMMA warp issues the block's K slices half by half + - The math warpgroup drains and releases each half on its own (double-buffered TMEM loads), + so the MMAs of the next block's first half overlap the epilogue of the current block's + second half + - Per-element epilogue math and the FMA order of every row are unchanged + +Early KV loads (early_kv): + - Launched with programmatic dependent launch + - Before griddepcontrol.wait the TMA warps read the schedule, context lengths and block + table and load up to num_kv_stages KV blocks per group that end before the request's + newest next_n tokens + - Q, weights and every later KV block are loaded after the wait; nothing is written before it + - The caller guarantees that the grids this launch depends on write no KV row older than + the newest next_n of each request, and no block table, context length or schedule entry + +Dynamic shape support: + - Model-constant dims (block_kv, head_dim, N, per_token) remain static for codegen + - Runtime-varying dims (batch_size, num_phys_blocks, max_ctx, max_blocks_per_seq, + num_ctas) are marked dynamic via mark_compact_shape_dynamic + - Allows JIT cache reuse across different batch sizes / sequence lengths + +Epilogue dtype flows (--acc_dtype / --epi_dtype): + + Flow 1: --acc_dtype fp32 --epi_dtype fp16 + Q(FP8) x K(FP8) -> MMA acc(FP32) -> TMEM(FP32) + -> LDTM -> Reg(FP32) -> cvt FP16 -> ReLU(FP16) + -> FMA(fma.rn.f16x2) with weights(FP16 from SMEM) -> partial sum(FP16) + -> x scale(FP32->FP16) -> cvt output_dtype -> store logits(output_dtype) + + Flow 2: --acc_dtype fp16 --epi_dtype fp16 + Q(FP8) x K(FP8) -> MMA acc(FP16) -> TMEM(FP16, pack_16b) + -> LDTM -> Reg(FP16) -> ReLU(FP16) + -> FMA(fma.rn.f16x2) with weights(FP16 from SMEM) -> partial sum(FP16) + -> x scale(FP32->FP16) -> cvt output_dtype -> store logits(output_dtype) + + --output_dtype: fp32 (default), fp16, bf16. Controls logits tensor dtype and final store conversion. + Default: --acc_dtype fp32 --epi_dtype fp32 --output_dtype fp32 +""" + +from typing import Tuple + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass import Float16, Int32 +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm, vector +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +# CuTe DSL CUDA 13 validates rounding modes as string literals. The string +# form is also accepted by older wrappers, so keep it version-independent. +_RND_RN = "rn" + + +@dsl_user_op +def pack_f16x2( + a: Float16, + b: Float16, + *, + loc=None, + ip=None, +) -> Int32: + f16_ty = Float16.mlir_type + i32_ty = Int32.mlir_type + vec2_f16 = ir.VectorType.get([2], f16_ty, loc=loc) + v = vector.from_elements( + vec2_f16, + (Float16(a).ir_value(loc=loc, ip=ip), Float16(b).ir_value(loc=loc, ip=ip)), + loc=loc, + ip=ip, + ) + return Int32(llvm.bitcast(i32_ty, v, loc=loc, ip=ip)) + + +@dsl_user_op +def unpack_f16x2( + packed: Int32, + *, + loc=None, + ip=None, +) -> Tuple[Float16, Float16]: + f16_ty = Float16.mlir_type + vec2_f16 = ir.VectorType.get([2], f16_ty, loc=loc) + v = llvm.bitcast(vec2_f16, Int32(packed).ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + r0 = Float16(vector.extract(v, dynamic_position=[], static_position=[0], loc=loc, ip=ip)) + r1 = Float16(vector.extract(v, dynamic_position=[], static_position=[1], loc=loc, ip=ip)) + return r0, r1 + + +@dsl_user_op +def fma_f16x2( + a: Int32, + b: Int32, + c: Int32, + *, + loc=None, + ip=None, +) -> Int32: + i32_ty = Int32.mlir_type + return Int32( + llvm.inline_asm( + i32_ty, + [ + Int32(a).ir_value(loc=loc, ip=ip), + Int32(b).ir_value(loc=loc, ip=ip), + Int32(c).ir_value(loc=loc, ip=ip), + ], + "fma.rn.f16x2 $0, $1, $2, $3;", + "=r,r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def max_f16x2( + a: Int32, + b: Int32, + *, + loc=None, + ip=None, +) -> Int32: + i32_ty = Int32.mlir_type + return Int32( + llvm.inline_asm( + i32_ty, + [Int32(a).ir_value(loc=loc, ip=ip), Int32(b).ir_value(loc=loc, ip=ip)], + "max.f16x2 $0, $1, $2;", + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def add_f16x2( + a: Int32, + b: Int32, + *, + loc=None, + ip=None, +) -> Int32: + i32_ty = Int32.mlir_type + return Int32( + llvm.inline_asm( + i32_ty, + [Int32(a).ir_value(loc=loc, ip=ip), Int32(b).ir_value(loc=loc, ip=ip)], + "add.f16x2 $0, $1, $2;", + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +class FP8MQALogitsKernel: + """FP8 paged MQA logits kernel for Blackwell (SM100). + + Each CTA processes a range of (q_idx, kv_split) pairs. + A split = 2 consecutive KV blocks within a sequence (one per warp group). + Q is shared between warp groups and reloaded when q_idx changes. + """ + + def __init__( + self, + block_kv: int = 128, + phys_block_kv: int = 128, + num_heads: int = 64, + head_dim: int = 128, + next_n: int = 1, + num_sms: int = 148, + remove_kv_wait_in_epilogue: bool = False, + early_tmem_copy: bool = False, + smem_subpartition_opt: bool = False, + max_kv_pipeline: bool = False, + max_umma_pipeline: bool = False, + num_epi_subtiles: int = 1, + epi_dtype=cutlass.Float32, + acc_dtype=cutlass.Float32, + output_dtype=cutlass.Float32, + weight_cache_limit: int | None = None, + split_n: bool = False, + early_kv: bool = False, + ): + self.block_kv = block_kv + self.phys_block_kv = phys_block_kv + self.num_blocks_per_mma = block_kv // phys_block_kv + assert block_kv % phys_block_kv == 0, ( + f"block_kv={block_kv} must be divisible by phys_block_kv={phys_block_kv}" + ) + assert self.num_blocks_per_mma <= 4, ( + f"num_blocks_per_mma={self.num_blocks_per_mma} exceeds max 4" + ) + self.remove_kv_wait_in_epilogue = remove_kv_wait_in_epilogue + self.early_tmem_copy = early_tmem_copy + self.smem_subpartition_opt = smem_subpartition_opt + self.num_heads = num_heads + if weight_cache_limit is not None and (weight_cache_limit <= 0 or weight_cache_limit % 4): + raise ValueError("weight_cache_limit must be a positive multiple of four") + self.weight_cache_limit = num_heads if weight_cache_limit is None else weight_cache_limit + self.head_dim = head_dim + self.next_n = next_n + self.N = next_n * num_heads + self.num_sms = num_sms + self.num_epi_subtiles = num_epi_subtiles + self.epi_dtype = epi_dtype + self.epi_bytes = 2 if epi_dtype == cutlass.Float16 else 4 + # sW stage stride padded to 128-byte SMEM alignment for TMA bulk copy. + # Without padding, e.g. fp16 + N=32 gives 64B per stage, so stage 1 + # at +64 would be misaligned (TMA requires 128-byte aligned SMEM dest). + w_stage_bytes = self.N * self.epi_bytes + self.w_stage_stride = ((w_stage_bytes + 127) // 128 * 128) // self.epi_bytes + self.output_dtype = output_dtype + if num_epi_subtiles > 1 and num_heads % num_epi_subtiles != 0: + raise ValueError("num_heads must be divisible by num_epi_subtiles") + if (num_heads // num_epi_subtiles) % 4 != 0: + raise ValueError( + "num_heads // num_epi_subtiles must be divisible by 4 (FMA unroll granularity)" + ) + if split_n and ( + next_n % 2 != 0 + or (self.N // 2) % 16 != 0 + or num_epi_subtiles != 1 + or epi_dtype != cutlass.Float32 + or acc_dtype != cutlass.Float32 + ): + raise ValueError( + "split_n needs an even next_n, N / 2 divisible by 16, one epilogue sub-tile " + "and an FP32 accumulator and epilogue" + ) + self.split_n = split_n + self.n_halves = 2 if split_n else 1 + self.early_kv = early_kv + self.num_groups = 2 + + self.num_math_threads = 256 + self.num_specialized_threads = 128 + self.threads_per_cta = 384 + self.num_math_warps = 8 + self.tma_warp_base = 8 + self.umma_warp_base = 10 + + self.num_q_stages = 3 # 3 stages for Q pipelining across batch sequences + + # TMEM: 512 columns total, each group needs N columns per UMMA stage + # max_umma_stages = 512 // (2 * N) + TMEM_COLS = 512 + if split_n: + # One TMEM stage per N half: 2 groups x 2 halves x N / 2 columns. + self.num_umma_stages = self.n_halves + elif max_umma_pipeline: + self.num_umma_stages = min(2, TMEM_COLS // (2 * self.N)) + else: + self.num_umma_stages = 1 + + if max_kv_pipeline: + smem_capacity = utils.get_smem_capacity_in_bytes() + # Reserve ~1 KB for barriers and misc + SMEM_BUDGET = smem_capacity - 1024 + # KV+Scale per stage (×2 groups): + # 2 * (block_kv * head_dim * 1B + block_kv * 4B) + kv_scale_per_stage = 2 * (block_kv * head_dim + block_kv * 4) + # Q+W per stage: Q is N * head_dim * 1B, W uses padded stride + qw_per_stage = self.N * head_dim + self.w_stage_stride * self.epi_bytes + qw_total = qw_per_stage * self.num_q_stages + self.num_kv_stages = (SMEM_BUDGET - qw_total) // kv_scale_per_stage + else: + self.num_kv_stages = 3 + + # Pad SMEM to push sW/sScales into sub-partition 1 (>= 128KB), + # avoiding sub-bank conflicts with UMMA reading sKV. + # Layout: barriers(~256B) | sKV_0 | sKV_1 | sQ | [pad] | sW | sScales + if self.smem_subpartition_opt: + BOUNDARY = 128 * 1024 + used = ( + 256 + + 2 * (block_kv * head_dim * self.num_kv_stages) + + self.N * head_dim * self.num_q_stages + ) + used = ((used + 127) // 128) * 128 + if used < BOUNDARY: + self.smem_pad_bytes = ((BOUNDARY - used + 1023) // 1024) * 1024 + else: + self.smem_pad_bytes = 0 + else: + self.smem_pad_bytes = 0 + + self.acc_dtype = acc_dtype + self.cta_group = tcgen05.CtaGroup.ONE + self.cluster_shape_mn = (1, 1) + self.mma_tiler_mn = (block_kv, self.N // self.n_halves) + + def _setup_mma(self, a_dtype, b_dtype, a_major, b_major): + self.a_dtype = a_dtype + self.b_dtype = b_dtype + self.a_major_mode = a_major + self.b_major_mode = b_major + + self.mma_tiler = (*self.mma_tiler_mn, 1) + tiled_mma = sm100_utils.make_trivial_tiled_mma( + a_dtype, + a_major, + b_major, + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) # 32 + + # Full-K: tile K = head_dim (128), 1 TMA per block + mma_inst_tile_k = self.head_dim // mma_inst_shape_k # 4 + full_k = mma_inst_shape_k * mma_inst_tile_k # 128 + self.mma_tiler = ( + self.mma_tiler_mn[0], + self.mma_tiler_mn[1], + full_k, + ) + + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.epi_tile = self.cta_tile_shape_mnk[:2] + + # KV SMEM: 3 stages per group, each stage holds full [128, 128] + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + a_dtype, + self.num_kv_stages, + ) + # Q SMEM: num_q_stages stages of [N, 128]; with split_n, Q stage s holds its two + # N halves as B stages 2s and 2s + 1. + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + b_dtype, + self.num_q_stages * self.n_halves, + ) + + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape) + self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake) + self.num_tmem_alloc_cols_total = ( + self.num_tmem_alloc_cols * self.num_groups * self.num_umma_stages + ) + + return tiled_mma + + def _q_b_stage(self, q_stage, hh): + """B (Q) SMEM stage of N half hh of Q stage q_stage.""" + if self.n_halves == 1: + return q_stage + return q_stage * self.n_halves + hh + + def _copy_q(self, q_copy, q_bar, q_stage, q_idx): + """TMA copies of request q_idx's Q (one box per N half) and weights into Q stage q_stage.""" + tma_atom_b, tBgB_h, tBsB, b_mcast_mask, tma_atom_w, tWgW, tWsW = q_copy + for hh in range(self.n_halves): + cute.copy( + tma_atom_b, + tBgB_h[hh][(None, 0, q_idx)], + tBsB[(None, self._q_b_stage(q_stage, hh))], + tma_bar_ptr=q_bar, + mcast_mask=b_mcast_mask, + ) + cute.copy( + tma_atom_w, + tWgW[(None, q_idx)], + tWsW[(None, q_stage)], + tma_bar_ptr=q_bar, + ) + + def _row_sum(self, acc_vec, t, w_cache, sW, q_stage, num_w_in_reg): + """sum_h w[t, h] * relu(acc[h]) of query row t, in the FP32 epilogue's FMA order.""" + zero = cutlass.Float32(0.0) + s0x, s0y, s1x, s1y = zero, zero, zero, zero + for h in range(0, self.num_heads, 4): + a0 = cutlass.max(acc_vec[h], zero) + a1 = cutlass.max(acc_vec[h + 1], zero) + a2 = cutlass.max(acc_vec[h + 2], zero) + a3 = cutlass.max(acc_vec[h + 3], zero) + if h < num_w_in_reg: + r0 = t * num_w_in_reg + h + w0, w1, w2, w3 = w_cache[r0], w_cache[r0 + 1], w_cache[r0 + 2], w_cache[r0 + 3] + else: + c0 = t * self.num_heads + h + w0 = sW[(c0, q_stage)] + w1 = sW[(c0 + 1, q_stage)] + w2 = sW[(c0 + 2, q_stage)] + w3 = sW[(c0 + 3, q_stage)] + s0x, s0y = cute.arch.fma_packed_f32x2((a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN) + s1x, s1y = cute.arch.fma_packed_f32x2((a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN) + return s0x + s0y + s1x + s1y + + @cute.jit + def __call__( + self, + kv_fused: cute.Tensor, # Fused KV: [num_phys_blocks, block_bytes] FP8 + b: cute.Tensor, # Q: [N, head_dim, batch_size] + weights: cute.Tensor, # [N, batch_size] (transposed for TMA) + logits: cute.Tensor, # [batch_size * next_n, max_context_len] + block_table: cute.Tensor, # [batch_size, max_blocks_per_seq] + context_lens: cute.Tensor, # [batch_size] + schedule_meta: cute.Tensor, # [num_sms+1, 2] int32 + num_phys_blocks: cutlass.Int32, + batch_size: cutlass.Int32, + stream: cuda.CUstream, + ): + # Derive KV and Scale views from fused buffer using CuTE ops. + # Fused layout per physical block: [KV data (phys_block_kv*head_dim)] [Scales (phys_block_kv*4)] + phys_block_kv = self.phys_block_kv + scale_offset_elems = phys_block_kv * self.head_dim # in FP8 elements + + # Recast fused buffer to FP8 (same 1-byte elements, needed for MMA type inference) + kv_fp8 = cute.recast_tensor(kv_fused, cutlass.Float8E4M3FN) + + # Q (b) was passed as uint8 to work around DLPack's lack of float8 support; + # recast back to FP8 so MMA type inference and TMA descriptors are correct. + b = cute.recast_tensor(b, cutlass.Float8E4M3FN) + + # Read the real per-block stride (bytes = FP8 elements) from the input. + # When KV is the indexer K-cache pool view, the pool is laid out as + # [num_blocks, num_layers, kvFactor, blockSize], so dim-0 stride = + # num_layers * kvFactor * phys_block_bytes (not phys_block_bytes). + # Using the input stride keeps both contiguous test path and the + # strided prod path correct. FP8 = 1 byte per element, so the byte + # stride is the same as the element stride after recast. + kv_block_stride = kv_fused.layout.stride[0] + + # KV view: [phys_block_kv, head_dim, num_phys_blocks] FP8 + # Each TMA loads one physical block; multiple TMAs fill a compute tile. + kv_layout = cute.make_layout( + (phys_block_kv, self.head_dim, num_phys_blocks), + stride=(self.head_dim, 1, kv_block_stride), + ) + a = cute.make_tensor(kv_fp8.iterator, kv_layout) + + # Scale view: offset pointer to scale region, recast FP8 → Float32 + # [phys_block_kv, num_phys_blocks] float32 (after recast) + scale_fp8_layout = cute.make_layout( + (phys_block_kv * 4, num_phys_blocks), + stride=(1, kv_block_stride), + ) + scale_fp8 = cute.make_tensor(kv_fp8.iterator + scale_offset_elems, scale_fp8_layout) + scales = cute.recast_tensor(scale_fp8, cutlass.Float32) + + a_dtype = a.element_type + b_dtype = b.element_type + a_major = utils.LayoutEnum.from_tensor(a).mma_major_mode() + b_major = utils.LayoutEnum.ROW_MAJOR.mma_major_mode() + + tiled_mma = self._setup_mma(a_dtype, b_dtype, a_major, b_major) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # TMA for KV (A) — fmha_decode_paged pattern. + # Build a TMA SMEM layout via tiled_divide on the full compute-tile + # layout, then select to drop trivial K dim. Atom uses mode [0] as + # single-tile SMEM layout and (phys, head) as cta_tiler. + tma_load_op = cpasync.CopyBulkTensorTileG2SOp() + self.a_tma_view_layout = sm100_utils.make_smem_layout( + tcgen05.OperandMajorMode.K, + (self.block_kv, self.head_dim), + a_dtype, + self.num_kv_stages, + ) + self.a_tma_view_layout = cute.tiled_divide( + self.a_tma_view_layout, (self.phys_block_kv, self.head_dim) + ) + # ((tile_M, tile_K), rest_M, rest_K, stages) → drop trivial rest_K + self.a_tma_view_layout = cute.select(self.a_tma_view_layout, mode=[0, 1, 3]) + # ((tile_M, tile_K), rest_M=num_sub_blocks, stages) + tma_atom_a, tma_tensor_a = cpasync.make_tiled_tma_atom( + tma_load_op, + a, + self.a_tma_view_layout[0], # atom SMEM = single-tile (mode 0) + (self.phys_block_kv, self.head_dim), + ) + + # TMA for Q (B) — full K=128, L dim = batch_size (unchanged) + b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA for Weights — [N, batch_size], tile [N], L=batch_size + self.w_smem_layout_staged = cute.make_layout( + (self.N, self.num_q_stages), + stride=(1, self.w_stage_stride), + ) + w_smem_per_stage = cute.select(self.w_smem_layout_staged, mode=[0]) + tma_atom_w, tma_tensor_w = cpasync.make_tiled_tma_atom( + tma_load_op, + weights, + w_smem_per_stage, + self.w_smem_layout_staged.shape[:1], + ) + + # TMA for Scales — [phys_block_kv, num_phys_blocks], tile [phys_block_kv] + # SMEM holds compute_block_kv scales per stage; filled by + # num_blocks_per_mma sub-block TMAs at consecutive offsets. + self.s_smem_layout_staged = cute.make_layout((self.block_kv, self.num_kv_stages)) + s_smem_per_subblock = cute.make_layout((phys_block_kv,)) + tma_atom_s, tma_tensor_s = cpasync.make_tiled_tma_atom( + tma_load_op, + scales, + s_smem_per_subblock, + (phys_block_kv,), + ) + + b_copy_size = cute.size_in_bytes(b_dtype, b_smem_layout) + w_copy_size = self.N * self.epi_bytes + # Per sub-block: phys_block_kv * head_dim (KV) + phys_block_kv * 4 (scales) + kv_tma_bytes_per_subblock = phys_block_kv * self.head_dim + scale_tma_bytes_per_subblock = phys_block_kv * 4 + # Total per compute tile = num_blocks_per_mma sub-blocks + self.num_kv_scale_tma_bytes = self.num_blocks_per_mma * ( + kv_tma_bytes_per_subblock + scale_tma_bytes_per_subblock + ) + # Q + Weights share barrier (like DeepGEMM); one Q box per N half + self.num_q_tma_bytes = b_copy_size * atom_thr_size * self.n_halves + w_copy_size + + num_ctas = self.num_sms + + @cute.struct + class SharedStorage: + kv_mbar_0: cute.struct.MemRange[cutlass.Int64, self.num_kv_stages * 2] + kv_mbar_1: cute.struct.MemRange[cutlass.Int64, self.num_kv_stages * 2] + q_mbar: cute.struct.MemRange[cutlass.Int64, self.num_q_stages * 2] + umma_mbar_0: cute.struct.MemRange[cutlass.Int64, self.num_umma_stages * 2] + umma_mbar_1: cute.struct.MemRange[cutlass.Int64, self.num_umma_stages * 2] + tmem_holding_buf: cutlass.Int32 + + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_w, + tma_tensor_w, + tma_atom_s, + tma_tensor_s, + logits, + block_table, + context_lens, + schedule_meta, + batch_size, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.w_smem_layout_staged, + self.s_smem_layout_staged, + self.a_tma_view_layout, + self.epi_tile, + SharedStorage, + ).launch( + grid=(1, 1, num_ctas), + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + use_pdl=self.early_kv, + ) + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, # KV pool + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, # Q (L dim = batch_size) + tma_atom_w: cute.CopyAtom, + mW_tma: cute.Tensor, # Weights TMA coord tensor [N, batch_size] + tma_atom_s: cute.CopyAtom, + mS_tma: cute.Tensor, # Scales TMA coord tensor [block_kv, num_phys_blocks] + mLogits: cute.Tensor, # [batch_size * next_n, max_context_len] + mBlockTable: cute.Tensor, # [batch_size, max_blocks_per_seq] + mContextLens: cute.Tensor, # [batch_size] + mScheduleMeta: cute.Tensor, # [num_sms+1, 2] int32 + batch_size: cutlass.Int32, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + w_smem_layout_staged: cute.Layout, + s_smem_layout_staged: cute.Layout, + a_tma_view_layout: cute.ComposedLayout, + epi_tile: cute.Tile, + SharedStorage: cutlass.Constexpr, + ): + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + tidx, _, _ = cute.arch.thread_idx() + + # Warp roles (matches DeepGEMM SM100) + warpgroup_idx = warp_idx // 4 + is_math_warp = warp_idx < 8 + is_tma_warp_0 = warp_idx == 8 + is_tma_warp_1 = warp_idx == 9 + is_tma_warp = is_tma_warp_0 | is_tma_warp_1 + is_umma_warp_0 = warp_idx == 10 + is_umma_warp_1 = warp_idx == 11 + + # Early schedule metadata load: issue global loads ASAP so their + # ~200-cycle L2 latency overlaps with subsequent prologue setup + # (SMEM alloc, TMA partition, MMA fragment creation, etc.) + NUM_MATH_WG = 2 # kNumMathWarpGroups + NUM_BLOCKS_PER_MMA = self.num_blocks_per_mma + sm_idx = bidz + start_q = mScheduleMeta[(sm_idx, 0)] + start_kv_half = mScheduleMeta[(sm_idx, 1)] + end_q_idx = mScheduleMeta[(sm_idx + 1, 0)] + end_kv_half = mScheduleMeta[(sm_idx + 1, 1)] + # Early mContextLens load: overlap ~200-cycle L2 latency with the + # entire prologue setup (pipelines, SMEM alloc, TMA partition, etc.) + # Clamp to avoid OOB when start_q == batch_size (zero-work CTA sentinel). + # Note: zero-work CTAs get a stale current_num_kv (from the last batch + # element), but it is never used because has_work will be False. + start_q_clamped = min(start_q, batch_size - 1) + current_num_kv = (mContextLens[start_q_clamped] + self.block_kv - 1) // self.block_kv + if cutlass.const_expr(self.early_kv): + current_ctx_tok = mContextLens[start_q_clamped] + + if is_tma_warp: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + cpasync.prefetch_descriptor(tma_atom_w) + cpasync.prefetch_descriptor(tma_atom_s) + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + block_kv_val = self.block_kv + num_heads = self.num_heads + next_n = self.next_n + num_epi_subtiles = self.num_epi_subtiles + + # === Pipelines === + prod_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + + # Q pipeline: TMA producer → Math consumer (8 math warps) + # PipelineTmaAsync: consumer_release uses is_signalling_thread + # (lane 0 per warp). 8 math warps × 1 lane-0 = 8 arrives. + q_cons_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 8) + q_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.q_mbar.data_ptr(), + num_stages=self.num_q_stages, + producer_group=prod_group, + consumer_group=q_cons_group, + tx_count=self.num_q_tma_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=tidx, + defer_sync=True, + ) + q_prod_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_q_stages + ) + # Both Math WGs share the same pipeline state (advance in lockstep) + q_cons_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_q_stages + ) + # UMMA warps observe Q pipeline (wait only, no release) + # to ensure Q is in SMEM before GEMM. Critical for UMMA warp 1 + # since TMA warp 1 only loads KV1 (not Q). + q_cons_state_umma_0 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_q_stages + ) + q_cons_state_umma_1 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_q_stages + ) + + # Merged KV+Scale pipelines (per-group, 3 stages each) + # Like DeepGEMM: KV data and scales share one barrier. + # TMA loads both under one barrier. Math is consumer (releases). + # UMMA also waits on this barrier (for KV GEMM) but does NOT release. + math_warps_per_group = self.num_math_warps // 2 # 4 warps + kv_cons_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, math_warps_per_group) + kv_pipeline_0 = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.kv_mbar_0.data_ptr(), + num_stages=self.num_kv_stages, + producer_group=prod_group, + consumer_group=kv_cons_group, + tx_count=self.num_kv_scale_tma_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=tidx, + defer_sync=True, + ) + kv_pipeline_1 = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.kv_mbar_1.data_ptr(), + num_stages=self.num_kv_stages, + producer_group=prod_group, + consumer_group=kv_cons_group, + tx_count=self.num_kv_scale_tma_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=tidx, + defer_sync=True, + ) + + kv_prod_state_0 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_kv_stages + ) + kv_prod_state_1 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_kv_stages + ) + # UMMA consumer states (wait only, no release) + kv_cons_state_umma_0 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_kv_stages + ) + kv_cons_state_umma_1 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_kv_stages + ) + # Math consumer states (wait + release) + kv_cons_state_math_0 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_kv_stages + ) + kv_cons_state_math_1 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_kv_stages + ) + + # UMMA pipelines (per-group) + math_threads_per_group = self.num_math_threads // 2 + umma_pipeline_0 = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.umma_mbar_0.data_ptr(), + num_stages=self.num_umma_stages, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, math_threads_per_group), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + umma_pipeline_1 = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.umma_mbar_1.data_ptr(), + num_stages=self.num_umma_stages, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, math_threads_per_group), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + umma_prod_state_0 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_umma_stages + ) + umma_prod_state_1 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_umma_stages + ) + umma_cons_state_0 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_umma_stages + ) + umma_cons_state_1 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_umma_stages + ) + + # TMEM — only Math warps (8×32=256) + UMMA warps (2×32=64) = 320 threads + # TMA warps do NOT participate, so they can start TMA loads earlier. + # Math warp 0 is the allocator (like fp16_gemm_3's epilogue warp 0), + # because math warps are the last TMEM consumers (epilogue reads). + tmem_alloc_num_threads = 320 # 10 warps: warp 0-7 (math) + warp 10-11 (umma) + tmem_alloc_barrier = pipeline.NamedBarrier(barrier_id=1, num_threads=tmem_alloc_num_threads) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=0, # math warp 0 does alloc+free (last TMEM consumer) + is_two_cta=False, + ) + + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # SMEM allocation: per-group KV + shared Q + sKV_0 = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + sKV_1 = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + sQ = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + # Pad SMEM to push sW/sScales into sub-partition 1 (>= 128KB) + # to avoid sub-bank conflicts with UMMA reading sKV from + # sub-partition 0. + if cutlass.const_expr(self.smem_pad_bytes > 0): + _ = smem.allocate(self.smem_pad_bytes) + # Weights SMEM: [N, num_q_stages], shared Q barrier + sW = smem.allocate_tensor( + element_type=self.epi_dtype, + layout=w_smem_layout_staged, + byte_alignment=128, + ) + # Scales SMEM: [block_kv, num_kv_stages] float32, per group + sScales_0 = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=s_smem_layout_staged, + byte_alignment=128, + ) + sScales_1 = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=s_smem_layout_staged, + byte_alignment=128, + ) + + a_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + # Partition KV (A): fmha_decode_paged pattern. + # SMEM view is ((tile), num_sub_blocks, stages) — built in __call__. + # Use .outer (plain layout); swizzle is captured by sKV_0's iterator. + # GMEM: local_tile by (phys, head), then group first 2 modes into tile. + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + sKV_0_for_tma = cute.make_tensor(sKV_0.iterator, a_tma_view_layout.outer) + sKV_1_for_tma = cute.make_tensor(sKV_1.iterator, a_tma_view_layout.outer) + gA = cute.local_tile( + mA_mkl, + (self.phys_block_kv, self.head_dim), + coord=(None, None, None), + ) + tAsA_0, tAgA_0 = cpasync.tma_partition( + tma_atom_a, + 0, + cute.make_layout(1), + sKV_0_for_tma, + cute.group_modes(gA, 0, 2), + ) + tAsA_1, tAgA_1 = cpasync.tma_partition( + tma_atom_a, + 0, + cute.make_layout(1), + sKV_1_for_tma, + cute.group_modes(gA, 0, 2), + ) + + # Partition Q (B): shared SMEM, L dim = batch_size + gB_nkl = cute.local_tile( + mB_nkl, + cute.slice_(self.mma_tiler, (0, None, None)), + (None, None, None), + ) + tCgB = thr_mma.partition_B(gB_nkl) + b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sQ, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + # [tma, K, L] per N half + tBgB_h = [tBgB[(None, hh, None, None)] for hh in range(self.n_halves)] + + # Partition Weights: standalone TMA, [N, batch_size] → [N] per stage + w_cta_layout = cute.make_layout((1,)) + tWsW, tWgW = cpasync.tma_partition( + tma_atom_w, + 0, + w_cta_layout, + cute.group_modes(sW, 0, 1), + cute.group_modes(mW_tma, 0, 1), + ) + q_copy = (tma_atom_b, tBgB_h, tBsB, b_mcast_mask, tma_atom_w, tWgW, tWsW) + + # Partition Scales: explicit sub_blocks + stages dims. + # Layout (phys_block_kv, num_sub, stages) K-major with custom strides. + s_tma_view_layout = cute.make_layout( + (self.phys_block_kv, self.num_blocks_per_mma, self.num_kv_stages), + stride=(1, self.phys_block_kv, self.block_kv), + ) + sScales_0_for_tma = cute.make_tensor(sScales_0.iterator, s_tma_view_layout) + sScales_1_for_tma = cute.make_tensor(sScales_1.iterator, s_tma_view_layout) + # GMEM: local_tile by phys to match atom's tile size + gS = cute.local_tile(mS_tma, (self.phys_block_kv,), coord=(None, None)) + tSsS_0, tSgS_0 = cpasync.tma_partition( + tma_atom_s, + 0, + cute.make_layout(1), + sScales_0_for_tma, + gS, + ) + tSsS_1, tSgS_1 = cpasync.tma_partition( + tma_atom_s, + 0, + cute.make_layout(1), + sScales_1_for_tma, + gS, + ) + + # MMA fragments + tCrA_0 = tiled_mma.make_fragment_A(sKV_0) + tCrA_1 = tiled_mma.make_fragment_A(sKV_1) + tCrB = tiled_mma.make_fragment_B(sQ) # shared + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + + # Staged acc (fp16_gemm_3 pattern): append UMMA stage dim + # shape: (*acc_shape, STAGE) — dynamic index on last dim reduces rank + us = self.num_umma_stages + cols = self.num_tmem_alloc_cols + acc_shape_staged = cute.append(acc_shape, us) + tCtAcc_fake_staged = tiled_mma.make_fragment_C(acc_shape_staged) + + # TMEM layout info (allocation deferred to UMMA/Math warp branches) + cols_per_group = cols * us * (32 // self.acc_dtype.width) + num_tmem_alloc_cols_total = self.num_tmem_alloc_cols_total + + # Epilogue setup + c_layout = utils.LayoutEnum.ROW_MAJOR + epi_sub_mn = (epi_tile[0], num_heads // num_epi_subtiles) + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + c_layout, + self.acc_dtype, + self.acc_dtype, + epi_sub_mn, + use_2cta_instrs, + ) + + # ===== SCHEDULER: derive values from early-loaded schedule metadata ===== + end_kv_idx = end_kv_half * NUM_MATH_WG + + # Convert start to KV block units (matching DeepGEMM) + current_q_idx = start_q + current_kv_idx = start_kv_half * NUM_MATH_WG + + # ===== COMMON SCHEDULER STATE (before warp branches) ===== + # Each warp role independently maintains its own copy of these + # variables (like DeepGEMM where each role creates its own scheduler). + # Pre-fetch first task (current_num_kv loaded early above for latency hiding) + next_q_idx = current_q_idx + next_kv_idx = current_kv_idx + next_num_kv = current_num_kv + if cutlass.const_expr(self.early_kv): + # TMA warps: context length of next_q_idx, and the CTA's first request + next_ctx_tok = current_ctx_tok + first_q = current_q_idx + # Sentinel: no previous batch (matches DeepGEMM's q_idx = batch_size) + q_idx = batch_size + # While-loop termination flag (matches DeepGEMM's fetch_next_task pattern). + # True if this CTA has work assigned (start != end in schedule_meta). + has_work = (current_q_idx != end_q_idx) | (current_kv_idx != end_kv_idx) + + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # ===== WARP-SPECIALIZED EXECUTION ===== + + if is_tma_warp_0: + # TMA warp 0: loads Q (prefetch) + KV for group 0 + # Matches DeepGEMM's TMA warp with kv_group_idx == 0 + cute.arch.warpgroup_reg_dealloc(24) + lane_idx = tidx % 32 + + # Block table prefetch: 32 lanes cache block indices, + # distributed via shuffle. Each lane holds num_blocks_per_mma + # physical block indices per compute tile. (DeepGEMM L233-244) + cached_blks = [cutlass.Int32(0) for _ in range(NUM_BLOCKS_PER_MMA)] + kv_blk_ptr = cutlass.Int32(32) # force prefetch on first use + + if cutlass.const_expr(not self.early_kv): + # Prefetch first Q before loop (like DeepGEMM line 203-204) + q_pipeline.producer_acquire(q_prod_state) + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + self._copy_q(q_copy, q_bar, q_prod_state.index, next_q_idx) + q_prod_state.advance() + else: + # Q is loaded after griddepcontrol.wait. Until then, count the KV + # tiles loaded; the first iteration's Q prefetch waits for it too. + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + waited = cutlass.Int32(0) + pre_wait_tiles = cutlass.Int32(0) + ran = cutlass.Int32(0) + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + if cutlass.const_expr(self.early_kv): + if waited == 0: + # Wait before the first tile that holds one of the request's newest + # next_n tokens, before a KV stage would be reused, and at a request change. + must_wait = cutlass.Int32(0) + if (kv_idx + 1) * block_kv_val > next_ctx_tok - next_n: + must_wait = cutlass.Int32(1) + if pre_wait_tiles >= self.num_kv_stages: + must_wait = cutlass.Int32(1) + if q_idx_old != batch_size: + if q_idx != q_idx_old: + must_wait = cutlass.Int32(1) + if must_wait != 0: + cute.arch.griddepcontrol_wait() + waited = cutlass.Int32(1) + q_pipeline.producer_acquire(q_prod_state) + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + self._copy_q(q_copy, q_bar, q_prod_state.index, first_q) + q_prod_state.advance() + if q_idx_old != batch_size: + # The Q prefetch the first iteration deferred + deferred_q = first_q + 1 + if deferred_q < end_q_idx: + q_pipeline.producer_acquire(q_prod_state) + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + self._copy_q(q_copy, q_bar, q_prod_state.index, deferred_q) + q_prod_state.advance() + elif deferred_q == end_q_idx: + if end_kv_idx > 0: + q_pipeline.producer_acquire(q_prod_state) + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + self._copy_q(q_copy, q_bar, q_prod_state.index, deferred_q) + q_prod_state.advance() + + # Q prefetch: when batch changes, load Q for NEXT batch + if q_idx != q_idx_old: + kv_blk_ptr = cutlass.Int32(32) # force re-prefetch + prefetch_next = q_idx + 1 + if cutlass.const_expr(self.early_kv): + if waited == 0: + # Deferred until the wait + prefetch_next = end_q_idx + 1 + if prefetch_next < end_q_idx: + q_pipeline.producer_acquire(q_prod_state) + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + self._copy_q(q_copy, q_bar, q_prod_state.index, prefetch_next) + q_prod_state.advance() + elif prefetch_next == end_q_idx: + if end_kv_idx > 0: + q_pipeline.producer_acquire(q_prod_state) + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + self._copy_q(q_copy, q_bar, q_prod_state.index, prefetch_next) + q_prod_state.advance() + + # Block table prefetch for group 0 (like DeepGEMM L233-241). + # Each lane loads num_blocks_per_mma physical block indices + # for one compute tile (kv_idx counts compute tiles). + if kv_blk_ptr == 32: + kv_blk_ptr = cutlass.Int32(0) + prefetch_kv = kv_idx + lane_idx * NUM_MATH_WG + if prefetch_kv < num_kv: + base_phys = prefetch_kv * NUM_BLOCKS_PER_MMA + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + cached_blks[i] = mBlockTable[(q_idx, base_phys + i)] + else: + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + cached_blks[i] = cutlass.Int32(0) + + # Get block indices via shuffle before barrier (like DeepGEMM L244) + phys_blks = [cutlass.Int32(0)] * NUM_BLOCKS_PER_MMA + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + phys_blks[i] = cute.arch.shuffle_sync(cached_blks[i], kv_blk_ptr) + kv_blk_ptr = kv_blk_ptr + 1 + + # Load KV + Scale for group 0: num_blocks_per_mma TMAs per tile. + kv_pipeline_0.producer_acquire(kv_prod_state_0) + bar = kv_pipeline_0.producer_get_barrier(kv_prod_state_0) + stage = kv_prod_state_0.index + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + cute.copy( + tma_atom_a, + tAgA_0[(None, 0, 0, phys_blks[i])], + tAsA_0[(None, i, stage)], + tma_bar_ptr=bar, + mcast_mask=a_mcast_mask, + ) + cute.copy( + tma_atom_s, + tSgS_0[(None, 0, phys_blks[i])], + tSsS_0[(None, i, stage)], + tma_bar_ptr=bar, + ) + kv_prod_state_0.advance() + if cutlass.const_expr(self.early_kv): + ran = cutlass.Int32(1) + if waited == 0: + pre_wait_tiles = pre_wait_tiles + 1 + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = (mContextLens[next_q_idx] + block_kv_val - 1) // block_kv_val + if cutlass.const_expr(self.early_kv): + next_ctx_tok = mContextLens[next_q_idx] + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + if cutlass.const_expr(self.early_kv): + if waited == 0: + # Every tile of this CTA was loaded before the wait (or it has none). + cute.arch.griddepcontrol_wait() + q_pipeline.producer_acquire(q_prod_state) + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + self._copy_q(q_copy, q_bar, q_prod_state.index, first_q) + q_prod_state.advance() + if ran != 0: + # The Q prefetch the first iteration deferred + deferred_q = first_q + 1 + if deferred_q < end_q_idx: + q_pipeline.producer_acquire(q_prod_state) + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + self._copy_q(q_copy, q_bar, q_prod_state.index, deferred_q) + q_prod_state.advance() + elif deferred_q == end_q_idx: + if end_kv_idx > 0: + q_pipeline.producer_acquire(q_prod_state) + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + self._copy_q(q_copy, q_bar, q_prod_state.index, deferred_q) + q_prod_state.advance() + + elif is_tma_warp_1: + # TMA warp 1: loads KV + Scale for group 1 only + # Matches DeepGEMM's TMA warp with kv_group_idx == 1 + cute.arch.warpgroup_reg_dealloc(24) + lane_idx = tidx % 32 + + # Block table prefetch for group 1 + cached_blks = [cutlass.Int32(0) for _ in range(NUM_BLOCKS_PER_MMA)] + kv_blk_ptr = cutlass.Int32(32) # force prefetch on first use + if cutlass.const_expr(self.early_kv): + waited = cutlass.Int32(0) + pre_wait_tiles = cutlass.Int32(0) + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + if cutlass.const_expr(self.early_kv): + if waited == 0: + # Same rule as TMA warp 0, for tile kv_idx + 1 + must_wait = cutlass.Int32(0) + if (kv_idx + 2) * block_kv_val > next_ctx_tok - next_n: + must_wait = cutlass.Int32(1) + if pre_wait_tiles >= self.num_kv_stages: + must_wait = cutlass.Int32(1) + if q_idx_old != batch_size: + if q_idx != q_idx_old: + must_wait = cutlass.Int32(1) + if must_wait != 0: + cute.arch.griddepcontrol_wait() + waited = cutlass.Int32(1) + + # New q_idx → force block table re-prefetch + if q_idx != q_idx_old: + kv_blk_ptr = cutlass.Int32(32) + + # Block table prefetch for group 1 (like DeepGEMM L233-241). + if kv_blk_ptr == 32: + kv_blk_ptr = cutlass.Int32(0) + prefetch_kv = kv_idx + 1 + lane_idx * NUM_MATH_WG + if prefetch_kv < num_kv: + base_phys = prefetch_kv * NUM_BLOCKS_PER_MMA + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + cached_blks[i] = mBlockTable[(q_idx, base_phys + i)] + else: + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + cached_blks[i] = cutlass.Int32(0) + + # Get block indices via shuffle before barrier (like DeepGEMM L244) + phys_blks = [cutlass.Int32(0)] * NUM_BLOCKS_PER_MMA + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + phys_blks[i] = cute.arch.shuffle_sync(cached_blks[i], kv_blk_ptr) + kv_blk_ptr = kv_blk_ptr + 1 + + # Load KV + Scale for group 1: num_blocks_per_mma TMAs per tile. + kv_pipeline_1.producer_acquire(kv_prod_state_1) + bar = kv_pipeline_1.producer_get_barrier(kv_prod_state_1) + stage = kv_prod_state_1.index + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + cute.copy( + tma_atom_a, + tAgA_1[(None, 0, 0, phys_blks[i])], + tAsA_1[(None, i, stage)], + tma_bar_ptr=bar, + mcast_mask=a_mcast_mask, + ) + cute.copy( + tma_atom_s, + tSgS_1[(None, 0, phys_blks[i])], + tSsS_1[(None, i, stage)], + tma_bar_ptr=bar, + ) + kv_prod_state_1.advance() + if cutlass.const_expr(self.early_kv): + if waited == 0: + pre_wait_tiles = pre_wait_tiles + 1 + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = (mContextLens[next_q_idx] + block_kv_val - 1) // block_kv_val + if cutlass.const_expr(self.early_kv): + next_ctx_tok = mContextLens[next_q_idx] + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + elif is_umma_warp_0: + # UMMA warp for group 0 + # Must wait on Q pipeline: TMA operations with different + # barriers are NOT visibility-ordered even within the same + # warp. KV0 barrier arriving does not guarantee Q SMEM + # writes are visible. + cute.arch.warpgroup_reg_dealloc(24) + + # TMEM: wait for math warp 0's allocation, retrieve pointer + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base_0 = cute.make_tensor(tmem_ptr, tCtAcc_fake_staged.layout) + tCtAcc_base_1 = cute.make_tensor(tmem_ptr + cols_per_group, tCtAcc_fake_staged.layout) + + if is_leader_cta: + num_k_blocks = cute.size(tCrA_0.shape[2]) + q_stage_0 = cutlass.Int32(0) + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + # Wait for Q pipeline when batch changes + if q_idx != q_idx_old: + if q_idx_old < batch_size: + q_cons_state_umma_0.advance() + q_pipeline.consumer_wait(q_cons_state_umma_0) + q_stage_0 = q_cons_state_umma_0.index + + # Process KV block for group 0 (kv_idx + 0) + # Unconditional UMMA (like DeepGEMM): OOB iterations + # compute on garbage data; results written to aligned + # padding region in logits buffer. + # Wait KV first, then TMEM empty (like DeepGEMM) + kv_pipeline_0.consumer_wait(kv_cons_state_umma_0) + # One accumulator stage per N half (a single one without split_n) + for hh in cutlass.range_constexpr(self.n_halves): + umma_pipeline_0.producer_acquire(umma_prod_state_0) + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + kv_stage = kv_cons_state_umma_0.index + q_b_stage = self._q_b_stage(q_stage_0, hh) + tCtAcc_0 = tCtAcc_base_0[(None, None, None, umma_prod_state_0.index)] + for k_block in cutlass.range_constexpr(num_k_blocks): + cute.gemm( + tiled_mma, + tCtAcc_0, + tCrA_0[None, None, k_block, kv_stage], + tCrB[None, None, k_block, q_b_stage], + tCtAcc_0, + ) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + if cutlass.const_expr(hh == self.n_halves - 1): + # No consumer_release here — Math WG 0 releases + kv_cons_state_umma_0.advance() + + umma_pipeline_0.producer_commit(umma_prod_state_0) + umma_prod_state_0.advance() + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = ( + mContextLens[next_q_idx] + block_kv_val - 1 + ) // block_kv_val + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + elif is_umma_warp_1: + # UMMA warp for group 1 + # Explicitly waits on Q pipeline — critical because TMA warp 1 + # only loads KV1, not Q. Without this wait, UMMA warp 1 can + # start GEMM before TMA warp 0 finishes loading Q into SMEM. + cute.arch.warpgroup_reg_dealloc(24) + + # TMEM: wait for umma_warp_0's allocation, retrieve pointer + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base_0 = cute.make_tensor(tmem_ptr, tCtAcc_fake_staged.layout) + tCtAcc_base_1 = cute.make_tensor(tmem_ptr + cols_per_group, tCtAcc_fake_staged.layout) + + if is_leader_cta: + num_k_blocks_1 = cute.size(tCrA_1.shape[2]) + q_stage_1 = cutlass.Int32(0) + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + # Wait for Q pipeline when batch changes + if q_idx != q_idx_old: + if q_idx_old < batch_size: + q_cons_state_umma_1.advance() + q_pipeline.consumer_wait(q_cons_state_umma_1) + q_stage_1 = q_cons_state_umma_1.index + + # Process KV block for group 1 (kv_idx + 1) + # Unconditional UMMA (like DeepGEMM) + # Wait KV first, then TMEM empty (like DeepGEMM) + kv_pipeline_1.consumer_wait(kv_cons_state_umma_1) + # One accumulator stage per N half (a single one without split_n) + for hh in cutlass.range_constexpr(self.n_halves): + umma_pipeline_1.producer_acquire(umma_prod_state_1) + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + kv_stage_1 = kv_cons_state_umma_1.index + q_b_stage = self._q_b_stage(q_stage_1, hh) + tCtAcc_1 = tCtAcc_base_1[(None, None, None, umma_prod_state_1.index)] + for k_block in cutlass.range_constexpr(num_k_blocks_1): + cute.gemm( + tiled_mma, + tCtAcc_1, + tCrA_1[None, None, k_block, kv_stage_1], + tCrB[None, None, k_block, q_b_stage], + tCtAcc_1, + ) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + if cutlass.const_expr(hh == self.n_halves - 1): + # No consumer_release here — Math WG 1 releases + kv_cons_state_umma_1.advance() + + umma_pipeline_1.producer_commit(umma_prod_state_1) + umma_prod_state_1.advance() + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = ( + mContextLens[next_q_idx] + block_kv_val - 1 + ) // block_kv_val + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + elif is_math_warp: + cute.arch.warpgroup_reg_alloc(240) + + # TMEM: math warp 0 is the allocator; all math warps wait + retrieve + tmem.allocate(num_tmem_alloc_cols_total) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base_0 = cute.make_tensor(tmem_ptr, tCtAcc_fake_staged.layout) + tCtAcc_base_1 = cute.make_tensor(tmem_ptr + cols_per_group, tCtAcc_fake_staged.layout) + + local_tidx = tidx % 128 + cC = cute.make_identity_tensor(epi_sub_mn) + + if cutlass.const_expr(self.split_n): + # Per KV block: wait for each N half's accumulator, drain its next_n / 2 rows + # (double-buffered TMEM loads) and release it; the KV stage is released once the + # MMAs of both halves are complete. + rows_per_half = next_n // self.n_halves + # Weight register cache as the FP32 epilogue below (next_n is even here) + MAX_NUM_W_IN_REG = 40 if next_n >= 4 else 52 + NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads, self.weight_cache_limit) + if warpgroup_idx == 0: + # Math WG 0: process group 0 + tAcc_ref = tCtAcc_base_0[(None, None, None, 0)][((None, None), 0, 0)] + tAcc_ref_epi = cute.flat_divide(tAcc_ref, epi_sub_mn) + tiled_copy_ref = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_ref_epi[(None, None, 0, 0)] + ) + tTR_cC = tiled_copy_ref.get_slice(local_tidx).partition_D(cC) + m_coord = tTR_cC[0][0] + tTR_rAcc = cute.make_fragment_like(tTR_cC, self.acc_dtype) + tTR_rAcc_next = cute.make_fragment_like(tTR_cC, self.acc_dtype) + w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype) + q_stage_local = cutlass.Int32(0) + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + # Q pipeline consumer: wait for Q+Weights SMEM + if q_idx != q_idx_old: + if q_idx_old < batch_size: + q_pipeline.consumer_release(q_cons_state) + q_cons_state.advance() + q_pipeline.consumer_wait(q_cons_state) + q_stage_local = q_cons_state.index + # Preload first NUM_W_IN_REG weights per slot + for t_i in cutlass.range_constexpr(next_n): + for w_j in cutlass.range_constexpr(NUM_W_IN_REG): + w_cache[t_i * NUM_W_IN_REG + w_j] = sW[ + (t_i * num_heads + w_j, q_stage_local) + ] + + # Process KV block for group 0 (kv_idx + 0) + kv_pos = kv_idx * block_kv_val + m_coord + kv_pipeline_0.consumer_wait(kv_cons_state_math_0) + scale_val = sScales_0[(m_coord, kv_cons_state_math_0.index)] + for hh in cutlass.range_constexpr(self.n_halves): + umma_pipeline_0.consumer_wait(umma_cons_state_0) + if cutlass.const_expr(hh == self.n_halves - 1): + # The MMAs of both halves are complete: the KV stage is free. + kv_pipeline_0.consumer_release(kv_cons_state_math_0) + kv_cons_state_math_0.advance() + tAcc_c = tCtAcc_base_0[(None, None, None, umma_cons_state_0.index)][ + ((None, None), 0, 0) + ] + tAcc_c_epi = cute.flat_divide(tAcc_c, epi_sub_mn) + tc = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_c_epi[(None, None, 0, 0)] + ) + tTR = tc.get_slice(local_tidx).partition_S(tAcc_c_epi) + cute.copy(tc, tTR[(None, None, None, 0, 0)], tTR_rAcc) + cute.arch.fence_view_async_tmem_load() + for j in cutlass.range_constexpr(rows_per_half): + t = hh * rows_per_half + j + cur = tTR_rAcc if j % 2 == 0 else tTR_rAcc_next + if cutlass.const_expr(j + 1 < rows_per_half): + # Next row's TMEM load into the other buffer + cute.copy( + tc, + tTR[(None, None, None, 0, j + 1)], + tTR_rAcc_next if j % 2 == 0 else tTR_rAcc, + ) + if cutlass.const_expr(j == rows_per_half - 1): + # Every TMEM load of this half has completed: release its stage. + umma_pipeline_0.consumer_release(umma_cons_state_0) + result_t = self._row_sum( + cur.load(), t, w_cache, sW, q_stage_local, NUM_W_IN_REG + ) + out_row = q_idx * next_n + t + mLogits[(out_row, kv_pos)] = self.output_dtype(result_t * scale_val) + if cutlass.const_expr(j + 1 < rows_per_half): + cute.arch.fence_view_async_tmem_load() + umma_cons_state_0.advance() + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = ( + mContextLens[next_q_idx] + block_kv_val - 1 + ) // block_kv_val + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + # Release last Q stage (WG 0) + if q_idx < batch_size: + q_pipeline.consumer_release(q_cons_state) + q_cons_state.advance() + else: + # Math WG 1: process group 1 + tAcc_ref = tCtAcc_base_1[(None, None, None, 0)][((None, None), 0, 0)] + tAcc_ref_epi = cute.flat_divide(tAcc_ref, epi_sub_mn) + tiled_copy_ref = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_ref_epi[(None, None, 0, 0)] + ) + tTR_cC = tiled_copy_ref.get_slice(local_tidx).partition_D(cC) + m_coord = tTR_cC[0][0] + tTR_rAcc = cute.make_fragment_like(tTR_cC, self.acc_dtype) + tTR_rAcc_next = cute.make_fragment_like(tTR_cC, self.acc_dtype) + w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype) + q_stage_local = cutlass.Int32(0) + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + # Q pipeline consumer: wait for Q+Weights SMEM + if q_idx != q_idx_old: + if q_idx_old < batch_size: + q_pipeline.consumer_release(q_cons_state) + q_cons_state.advance() + q_pipeline.consumer_wait(q_cons_state) + q_stage_local = q_cons_state.index + # Preload first NUM_W_IN_REG weights per slot + for t_i in cutlass.range_constexpr(next_n): + for w_j in cutlass.range_constexpr(NUM_W_IN_REG): + w_cache[t_i * NUM_W_IN_REG + w_j] = sW[ + (t_i * num_heads + w_j, q_stage_local) + ] + + # Process KV block for group 1 (kv_idx + 1) + kv_pos = (kv_idx + 1) * block_kv_val + m_coord + kv_pipeline_1.consumer_wait(kv_cons_state_math_1) + scale_val = sScales_1[(m_coord, kv_cons_state_math_1.index)] + for hh in cutlass.range_constexpr(self.n_halves): + umma_pipeline_1.consumer_wait(umma_cons_state_1) + if cutlass.const_expr(hh == self.n_halves - 1): + # The MMAs of both halves are complete: the KV stage is free. + kv_pipeline_1.consumer_release(kv_cons_state_math_1) + kv_cons_state_math_1.advance() + tAcc_c = tCtAcc_base_1[(None, None, None, umma_cons_state_1.index)][ + ((None, None), 0, 0) + ] + tAcc_c_epi = cute.flat_divide(tAcc_c, epi_sub_mn) + tc = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_c_epi[(None, None, 0, 0)] + ) + tTR = tc.get_slice(local_tidx).partition_S(tAcc_c_epi) + cute.copy(tc, tTR[(None, None, None, 0, 0)], tTR_rAcc) + cute.arch.fence_view_async_tmem_load() + for j in cutlass.range_constexpr(rows_per_half): + t = hh * rows_per_half + j + cur = tTR_rAcc if j % 2 == 0 else tTR_rAcc_next + if cutlass.const_expr(j + 1 < rows_per_half): + # Next row's TMEM load into the other buffer + cute.copy( + tc, + tTR[(None, None, None, 0, j + 1)], + tTR_rAcc_next if j % 2 == 0 else tTR_rAcc, + ) + if cutlass.const_expr(j == rows_per_half - 1): + # Every TMEM load of this half has completed: release its stage. + umma_pipeline_1.consumer_release(umma_cons_state_1) + result_t = self._row_sum( + cur.load(), t, w_cache, sW, q_stage_local, NUM_W_IN_REG + ) + out_row = q_idx * next_n + t + mLogits[(out_row, kv_pos)] = self.output_dtype(result_t * scale_val) + if cutlass.const_expr(j + 1 < rows_per_half): + cute.arch.fence_view_async_tmem_load() + umma_cons_state_1.advance() + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = ( + mContextLens[next_q_idx] + block_kv_val - 1 + ) // block_kv_val + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + # Release last Q stage (WG 1) + if q_idx < batch_size: + q_pipeline.consumer_release(q_cons_state) + q_cons_state.advance() + else: + if warpgroup_idx == 0: + # Math WG 0: process group 0 + # Reference setup (stage 0) for m_coord + # flat_divide by sub-tile to get sub-tile partitions + tAcc_0_ref = tCtAcc_base_0[(None, None, None, 0)][((None, None), 0, 0)] + tAcc_0_ref_epi = cute.flat_divide(tAcc_0_ref, epi_sub_mn) + tiled_copy_ref_0 = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_0_ref_epi[(None, None, 0, 0)] + ) + thr_copy_ref_0 = tiled_copy_ref_0.get_slice(local_tidx) + tTR_cC = thr_copy_ref_0.partition_D(cC) + m_coord = tTR_cC[0][0] + + tTR_rAcc = cute.make_fragment_like(tTR_cC, self.acc_dtype) + + # Weight register cache: only first NUM_W_IN_REG + # per next_n slot (like DeepGEMM min(52, kNumHeads)). + # Remaining weights read from SMEM in epilogue. + # FP16 weights use half the regs, so we can fit + # all heads for next_n <= 3. + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + MAX_NUM_W_IN_REG = 64 if next_n <= 3 else 48 + else: + MAX_NUM_W_IN_REG = 64 if next_n == 1 else 40 if next_n >= 4 else 52 + NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads, self.weight_cache_limit) + w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype) + q_stage_local = cutlass.Int32(0) + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + # Q pipeline consumer: wait for Q+Weights SMEM + if q_idx != q_idx_old: + if q_idx_old < batch_size: + q_pipeline.consumer_release(q_cons_state) + q_cons_state.advance() + q_pipeline.consumer_wait(q_cons_state) + q_stage_local = q_cons_state.index + # Preload first NUM_W_IN_REG weights per slot + for t_i in cutlass.range_constexpr(next_n): + for w_j in cutlass.range_constexpr(NUM_W_IN_REG): + w_cache[t_i * NUM_W_IN_REG + w_j] = sW[ + (t_i * num_heads + w_j, q_stage_local) + ] + + # Process KV block for group 0 (kv_idx + 0) + # Unconditional Math (like DeepGEMM): OOB results + # written to aligned padding region in logits buffer. + kv_pos = kv_idx * block_kv_val + m_coord + + if cutlass.const_expr(self.remove_kv_wait_in_epilogue): + # Skip KV wait, rely on UMMA barrier's + # transitive visibility. + umma_pipeline_0.consumer_wait(umma_cons_state_0) + else: + # Default: wait KV first to overlap lds with + # UMMA computation. + kv_pipeline_0.consumer_wait(kv_cons_state_math_0) + sc_stage_0 = kv_cons_state_math_0.index + scale_val = sScales_0[(m_coord, sc_stage_0)] + umma_pipeline_0.consumer_wait(umma_cons_state_0) + + # --- TMEM sub-tile setup --- + # flat_divide accumulator by sub-tile shape; + # partition once, then loop over sub-tiles. + tCtAcc_c0 = tCtAcc_base_0[(None, None, None, umma_cons_state_0.index)] + tAcc_c0 = tCtAcc_c0[((None, None), 0, 0)] + tAcc_c0_epi = cute.flat_divide(tAcc_c0, epi_sub_mn) + tc_0 = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_c0_epi[(None, None, 0, 0)] + ) + tr_0 = tc_0.get_slice(local_tidx) + tTR_0 = tr_0.partition_S(tAcc_c0_epi) + + # --- First sub-tile LDTM + KV release --- + if cutlass.const_expr(self.early_tmem_copy): + # Issue first sub-tile LDTM early + cute.copy(tc_0, tTR_0[(None, None, None, 0, 0)], tTR_rAcc) + # Scale LDS + KV release fill latency + if cutlass.const_expr(self.remove_kv_wait_in_epilogue): + sc_stage_0 = kv_cons_state_math_0.index + scale_val = sScales_0[(m_coord, sc_stage_0)] + kv_pipeline_0.consumer_release(kv_cons_state_math_0) + kv_cons_state_math_0.advance() + cute.arch.fence_view_async_tmem_load() + else: + # Default: scale LDS + KV release first + if cutlass.const_expr(self.remove_kv_wait_in_epilogue): + sc_stage_0 = kv_cons_state_math_0.index + scale_val = sScales_0[(m_coord, sc_stage_0)] + kv_pipeline_0.consumer_release(kv_cons_state_math_0) + kv_cons_state_math_0.advance() + cute.copy(tc_0, tTR_0[(None, None, None, 0, 0)], tTR_rAcc) + cute.arch.fence_view_async_tmem_load() + + # --- Sub-tile compute loop --- + # Each sub-tile: LDTM.xN → fence → load → + # ReLU+FMA. Breaks FMA chain (16→4 per chunk) + # and interleaves LDTM with FP32 compute to + # reduce ShadowPipeThrottle. + # Sub-tiles are within each t-slot (num_heads + # // num_epi_subtiles wide). flat_divide yields + # next_n * num_epi_subtiles sub-tiles total; + # global index = t * num_epi_subtiles + i. + subtile_n = num_heads // num_epi_subtiles + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + packed_zero = pack_f16x2(Float16(0.0), Float16(0.0)) + for t in cutlass.range_constexpr(next_n): + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + ps0 = packed_zero + ps1 = packed_zero + else: + s0x = cutlass.Float32(0.0) + s0y = cutlass.Float32(0.0) + s1x = cutlass.Float32(0.0) + s1y = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(num_epi_subtiles): + # LDTM for sub-tiles 1..N-1 + # (sub-tile 0 handled above) + if t > 0 or i > 0: + cute.copy( + tc_0, + tTR_0[(None, None, None, 0, t * num_epi_subtiles + i)], + tTR_rAcc, + ) + cute.arch.fence_view_async_tmem_load() + # Release UMMA after last LDTM+fence + if t == next_n - 1 and i == num_epi_subtiles - 1: + umma_pipeline_0.consumer_release(umma_cons_state_0) + umma_cons_state_0.advance() + acc_vec = tTR_rAcc.load() + # Reg-path: weights from registers + reg_h_end = min(subtile_n, max(0, NUM_W_IN_REG - i * subtile_n)) + for h in cutlass.range_constexpr(0, reg_h_end, 4): + n0 = h + h_g = i * subtile_n + h + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + pa01 = pack_f16x2( + Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1]) + ) + pa23 = pack_f16x2( + Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3]) + ) + pa01 = max_f16x2(pa01, packed_zero) + pa23 = max_f16x2(pa23, packed_zero) + r0 = t * NUM_W_IN_REG + h_g + pw01 = pack_f16x2(w_cache[r0], w_cache[r0 + 1]) + pw23 = pack_f16x2(w_cache[r0 + 2], w_cache[r0 + 3]) + ps0 = fma_f16x2(pa01, pw01, ps0) + ps1 = fma_f16x2(pa23, pw23, ps1) + else: + a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0)) + a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0)) + a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0)) + a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0)) + r0 = t * NUM_W_IN_REG + h_g + w0 = w_cache[r0] + w1 = w_cache[r0 + 1] + w2 = w_cache[r0 + 2] + w3 = w_cache[r0 + 3] + s0x, s0y = cute.arch.fma_packed_f32x2( + (a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN + ) + s1x, s1y = cute.arch.fma_packed_f32x2( + (a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN + ) + # SMEM-path: weights from shared mem + smem_h_start = max(0, NUM_W_IN_REG - i * subtile_n) + for h in cutlass.range_constexpr(smem_h_start, subtile_n, 4): + n0 = h + h_g = i * subtile_n + h + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + pa01 = pack_f16x2( + Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1]) + ) + pa23 = pack_f16x2( + Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3]) + ) + pa01 = max_f16x2(pa01, packed_zero) + pa23 = max_f16x2(pa23, packed_zero) + pw01 = pack_f16x2( + sW[(t * num_heads + h_g, q_stage_local)], + sW[(t * num_heads + h_g + 1, q_stage_local)], + ) + pw23 = pack_f16x2( + sW[(t * num_heads + h_g + 2, q_stage_local)], + sW[(t * num_heads + h_g + 3, q_stage_local)], + ) + ps0 = fma_f16x2(pa01, pw01, ps0) + ps1 = fma_f16x2(pa23, pw23, ps1) + else: + a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0)) + a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0)) + a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0)) + a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0)) + w0 = sW[(t * num_heads + h_g, q_stage_local)] + w1 = sW[(t * num_heads + h_g + 1, q_stage_local)] + w2 = sW[(t * num_heads + h_g + 2, q_stage_local)] + w3 = sW[(t * num_heads + h_g + 3, q_stage_local)] + s0x, s0y = cute.arch.fma_packed_f32x2( + (a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN + ) + s1x, s1y = cute.arch.fma_packed_f32x2( + (a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN + ) + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + ps_sum = add_f16x2(ps0, ps1) + sum_lo, sum_hi = unpack_f16x2(ps_sum) + result_t = sum_lo + sum_hi + else: + result_t = s0x + s0y + s1x + s1y + out_row = q_idx * next_n + t + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + mLogits[(out_row, kv_pos)] = self.output_dtype( + result_t * Float16(scale_val) + ) + else: + mLogits[(out_row, kv_pos)] = self.output_dtype(result_t * scale_val) + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = ( + mContextLens[next_q_idx] + block_kv_val - 1 + ) // block_kv_val + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + # Release last Q stage (WG 0) + if q_idx < batch_size: + q_pipeline.consumer_release(q_cons_state) + q_cons_state.advance() + + else: + # Math WG 1: process group 1 + tAcc_1_ref = tCtAcc_base_1[(None, None, None, 0)][((None, None), 0, 0)] + tAcc_1_ref_epi = cute.flat_divide(tAcc_1_ref, epi_sub_mn) + tiled_copy_ref_1 = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_1_ref_epi[(None, None, 0, 0)] + ) + thr_copy_ref_1 = tiled_copy_ref_1.get_slice(local_tidx) + tTR_cC = thr_copy_ref_1.partition_D(cC) + m_coord = tTR_cC[0][0] + + tTR_rAcc = cute.make_fragment_like(tTR_cC, self.acc_dtype) + + # Weight register cache: only first NUM_W_IN_REG + # per next_n slot (like DeepGEMM min(52, kNumHeads)). + # FP16 weights use half the regs, so we can fit + # all heads for next_n <= 3. + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + MAX_NUM_W_IN_REG = 64 if next_n <= 3 else 48 + else: + MAX_NUM_W_IN_REG = 64 if next_n == 1 else 40 if next_n >= 4 else 52 + NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads, self.weight_cache_limit) + w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype) + q_stage_local = cutlass.Int32(0) + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + # Q pipeline consumer: wait for Q+Weights SMEM + if q_idx != q_idx_old: + if q_idx_old < batch_size: + q_pipeline.consumer_release(q_cons_state) + q_cons_state.advance() + q_pipeline.consumer_wait(q_cons_state) + q_stage_local = q_cons_state.index + # Preload first NUM_W_IN_REG weights per slot + for t_i in cutlass.range_constexpr(next_n): + for w_j in cutlass.range_constexpr(NUM_W_IN_REG): + w_cache[t_i * NUM_W_IN_REG + w_j] = sW[ + (t_i * num_heads + w_j, q_stage_local) + ] + + # Process KV block for group 1 (kv_idx + 1) + # Unconditional Math (like DeepGEMM) + kv_idx_1 = kv_idx + 1 + + kv_pos = kv_idx_1 * block_kv_val + m_coord + + if cutlass.const_expr(self.remove_kv_wait_in_epilogue): + umma_pipeline_1.consumer_wait(umma_cons_state_1) + else: + kv_pipeline_1.consumer_wait(kv_cons_state_math_1) + sc_stage_1 = kv_cons_state_math_1.index + scale_val = sScales_1[(m_coord, sc_stage_1)] + umma_pipeline_1.consumer_wait(umma_cons_state_1) + + # --- TMEM sub-tile setup (WG1) --- + tCtAcc_c1 = tCtAcc_base_1[(None, None, None, umma_cons_state_1.index)] + tAcc_c1 = tCtAcc_c1[((None, None), 0, 0)] + tAcc_c1_epi = cute.flat_divide(tAcc_c1, epi_sub_mn) + tc_1 = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_c1_epi[(None, None, 0, 0)] + ) + tr_1 = tc_1.get_slice(local_tidx) + tTR_1 = tr_1.partition_S(tAcc_c1_epi) + + # --- First sub-tile LDTM + KV release (WG1) --- + if cutlass.const_expr(self.early_tmem_copy): + cute.copy(tc_1, tTR_1[(None, None, None, 0, 0)], tTR_rAcc) + if cutlass.const_expr(self.remove_kv_wait_in_epilogue): + sc_stage_1 = kv_cons_state_math_1.index + scale_val = sScales_1[(m_coord, sc_stage_1)] + kv_pipeline_1.consumer_release(kv_cons_state_math_1) + kv_cons_state_math_1.advance() + cute.arch.fence_view_async_tmem_load() + else: + if cutlass.const_expr(self.remove_kv_wait_in_epilogue): + sc_stage_1 = kv_cons_state_math_1.index + scale_val = sScales_1[(m_coord, sc_stage_1)] + kv_pipeline_1.consumer_release(kv_cons_state_math_1) + kv_cons_state_math_1.advance() + cute.copy(tc_1, tTR_1[(None, None, None, 0, 0)], tTR_rAcc) + cute.arch.fence_view_async_tmem_load() + + # --- Sub-tile compute loop (WG1) --- + subtile_n = num_heads // num_epi_subtiles + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + packed_zero = pack_f16x2(Float16(0.0), Float16(0.0)) + for t in cutlass.range_constexpr(next_n): + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + ps0 = packed_zero + ps1 = packed_zero + else: + s0x = cutlass.Float32(0.0) + s0y = cutlass.Float32(0.0) + s1x = cutlass.Float32(0.0) + s1y = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(num_epi_subtiles): + if t > 0 or i > 0: + cute.copy( + tc_1, + tTR_1[(None, None, None, 0, t * num_epi_subtiles + i)], + tTR_rAcc, + ) + cute.arch.fence_view_async_tmem_load() + if t == next_n - 1 and i == num_epi_subtiles - 1: + umma_pipeline_1.consumer_release(umma_cons_state_1) + umma_cons_state_1.advance() + acc_vec = tTR_rAcc.load() + # Reg-path + reg_h_end = min(subtile_n, max(0, NUM_W_IN_REG - i * subtile_n)) + for h in cutlass.range_constexpr(0, reg_h_end, 4): + n0 = h + h_g = i * subtile_n + h + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + pa01 = pack_f16x2( + Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1]) + ) + pa23 = pack_f16x2( + Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3]) + ) + pa01 = max_f16x2(pa01, packed_zero) + pa23 = max_f16x2(pa23, packed_zero) + r0 = t * NUM_W_IN_REG + h_g + pw01 = pack_f16x2(w_cache[r0], w_cache[r0 + 1]) + pw23 = pack_f16x2(w_cache[r0 + 2], w_cache[r0 + 3]) + ps0 = fma_f16x2(pa01, pw01, ps0) + ps1 = fma_f16x2(pa23, pw23, ps1) + else: + a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0)) + a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0)) + a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0)) + a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0)) + r0 = t * NUM_W_IN_REG + h_g + w0 = w_cache[r0] + w1 = w_cache[r0 + 1] + w2 = w_cache[r0 + 2] + w3 = w_cache[r0 + 3] + s0x, s0y = cute.arch.fma_packed_f32x2( + (a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN + ) + s1x, s1y = cute.arch.fma_packed_f32x2( + (a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN + ) + # SMEM-path + smem_h_start = max(0, NUM_W_IN_REG - i * subtile_n) + for h in cutlass.range_constexpr(smem_h_start, subtile_n, 4): + n0 = h + h_g = i * subtile_n + h + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + pa01 = pack_f16x2( + Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1]) + ) + pa23 = pack_f16x2( + Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3]) + ) + pa01 = max_f16x2(pa01, packed_zero) + pa23 = max_f16x2(pa23, packed_zero) + pw01 = pack_f16x2( + sW[(t * num_heads + h_g, q_stage_local)], + sW[(t * num_heads + h_g + 1, q_stage_local)], + ) + pw23 = pack_f16x2( + sW[(t * num_heads + h_g + 2, q_stage_local)], + sW[(t * num_heads + h_g + 3, q_stage_local)], + ) + ps0 = fma_f16x2(pa01, pw01, ps0) + ps1 = fma_f16x2(pa23, pw23, ps1) + else: + a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0)) + a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0)) + a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0)) + a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0)) + w0 = sW[(t * num_heads + h_g, q_stage_local)] + w1 = sW[(t * num_heads + h_g + 1, q_stage_local)] + w2 = sW[(t * num_heads + h_g + 2, q_stage_local)] + w3 = sW[(t * num_heads + h_g + 3, q_stage_local)] + s0x, s0y = cute.arch.fma_packed_f32x2( + (a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN + ) + s1x, s1y = cute.arch.fma_packed_f32x2( + (a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN + ) + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + ps_sum = add_f16x2(ps0, ps1) + sum_lo, sum_hi = unpack_f16x2(ps_sum) + result_t = sum_lo + sum_hi + else: + result_t = s0x + s0y + s1x + s1y + out_row = q_idx * next_n + t + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + mLogits[(out_row, kv_pos)] = self.output_dtype( + result_t * Float16(scale_val) + ) + else: + mLogits[(out_row, kv_pos)] = self.output_dtype(result_t * scale_val) + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = ( + mContextLens[next_q_idx] + block_kv_val - 1 + ) // block_kv_val + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + # Release last Q stage (WG 1) + if q_idx < batch_size: + q_pipeline.consumer_release(q_cons_state) + q_cons_state.advance() + + # TMEM dealloc: math warps are allocator + last consumer + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + else: + cute.arch.warpgroup_reg_dealloc(24) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/fp8_paged_mqa_logits.py new file mode 100644 index 000000000000..6ef531814cdb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/fp8_paged_mqa_logits.py @@ -0,0 +1,250 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Indexer FP8 MQA logits over paged keys (CuTe DSL), the whole verify window as one query atom.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._fp8_paged_mqa_logits import SPLIT_KV, num_sms, run_fp8_paged_mqa_logits + +#: The indexer K cache's page: 64 tokens of 128 FP8 bytes, then their 64 fp32 scales. +_PAGE = 64 +#: Columns of the c1 shard logits: the largest shard of the c1 gen server's indexer capacity +#: (1,048,596 keys) on 16 ranks. +_SHARD_CAPACITY = 65_664 + + +def _spec(context_lens, max_context_len=_SHARD_CAPACITY, **overrides) -> dict: + spec = dict( + next_n=8, + num_heads=32, + head_dim=128, + tokens_per_block=_PAGE, + context_lens=tuple(context_lens), + max_context_len=max_context_len, + early_kv=True, + data="int", + pool_layout="contiguous", + ) + spec.update(overrides) + return spec + + +class GlmFp8PagedMqaLogits(OpWrapper): + """``logits[b * next_n + t, j] = k_scale[j] * sum_h weights[b * next_n + t, h] * + relu(q[b, t, h] . k[j])`` for the keys ``j < context_lens[b]`` of request ``b``. + + The keys live in the indexer K cache: key ``j`` of request ``b`` is token ``j % page`` of page + ``block_table[b, j // page]``, whose bytes hold the page's FP8 rows and then their fp32 scales. + ``weights`` already carry the Q scales and the softmax scale (the indexer prologue's or the + wq_b epilogue's output). ``schedule_meta`` partitions the requests' 256-key splits over the + device's SMs (``paged_mqa_logits_schedule``); the kernel runs one persistent CTA per SM. + + A request's ``next_n`` query rows are one MMA atom of ``next_n * num_heads`` columns. With + ``early_kv`` (an even ``next_n`` whose half-atom is a whole number of 16-column UMMA atoms) the + N-split variant runs, launched with PDL: it loads the KV blocks that end before each request's + newest ``next_n`` keys before ``griddepcontrol.wait`` and everything else after it. Both + variants compute every element in the same order. + + ``logits`` is the caller's [B * next_n, >= round_up(max_context_len, 256)] fp32 buffer: the + kernel writes whole 256-key splits, so columns from ``context_lens[b]`` to the split's end hold + unspecified values. Returns the view ``logits[:, :max_context_len]``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="c1 key shard, 8 rows x 32 heads, early KV: a full 33,152-key shard", + spec=_spec([33_152]), + ), + Cell( + why="the request's last shard: 32,720 keys, a partial final page", + spec=_spec([32_720]), + ), + Cell( + why="an empty shard scores one key", + spec=_spec([1]), + ), + Cell( + why="the default kernel (early KV off) on the c1 shard", + spec=_spec([33_152], early_kv=False), + ), + Cell( + why="UE8M0-scaled keys and real-valued weights (accumulation order visible)", + spec=_spec([20_000], data="randn"), + ), + Cell( + why="the whole context on every rank (no key sharding): 200,000 keys of 1M", + spec=_spec([200_000], max_context_len=1_048_576), + ), + Cell( + why="cache pages interleaved with other layers' pages (strided page dim)", + spec=_spec([5_000], pool_layout="interleaved"), + ), + Cell( + why="four requests of 8 rows each (a larger batch's verify step)", + spec=_spec([3_000, 70_000, 129, 40_000], max_context_len=131_072), + ), + ) + + note = """ + Metadata consumed: none; every input is an argument. At c1 the target passes the step's Q + ([1, 8, 32, 128] E4M3) and scaled weights ([8, 32] fp32), the layer's indexer K-cache view + ([pages, 64, 1, 132] uint8), and, with the key-sharded top-k, the shard's context length, + block table and schedule (``hist_shard_meta``, ``paged_mqa_logits_schedule``) with + ``max_context_len`` = the shard capacity; without it, the request's indexer KV length, its + INDEX_KEY block-table row and the full schedule. + + The compiled kernel is cached per (page size, heads, head dim, next_n, SMs, weight-register + cache, early KV); the first call of a shape compiles (seconds), so it must run before CUDA + graph capture. + + PDL: only the early-KV variant is launched with the PDL attribute. It requires that the + kernels it waits for write no key older than each request's newest ``next_n`` and no block + table, context length or schedule entry. + + Not certified: FP16 epilogues, next_n other than 8, head counts other than 32, pages other + than 64 tokens. + """ + + def __call__( + self, + q: torch.Tensor, + kv_cache: torch.Tensor, + weights: torch.Tensor, + context_lens: torch.Tensor, + block_table: torch.Tensor, + schedule_meta: torch.Tensor, + logits: torch.Tensor, + max_context_len: int, + early_kv: bool, + ) -> torch.Tensor: + return run_fp8_paged_mqa_logits( + q, + kv_cache, + weights, + context_lens, + block_table, + schedule_meta, + logits, + max_context_len, + early_kv, + ) + + def reference( + self, + q: torch.Tensor, + kv_cache: torch.Tensor, + weights: torch.Tensor, + context_lens: torch.Tensor, + block_table: torch.Tensor, + schedule_meta: torch.Tensor, + logits: torch.Tensor, + max_context_len: int, + early_kv: bool, + ) -> torch.Tensor: + """[B * next_n, max_context_len] fp32, fp64-accumulated and rounded once; columns past a + request's context are 0. The schedule, the output buffer and the variant change where and + when the values are computed, not what they are.""" + batch, next_n, num_heads, head_dim = q.shape + page = kv_cache.shape[1] + out = torch.zeros(batch * next_n, max_context_len, dtype=torch.float32, device=q.device) + page_bytes = kv_cache.reshape(kv_cache.shape[0], -1) + for b, ctx in enumerate(context_lens.tolist()): + num_pages = -(-ctx // page) + pages = page_bytes[block_table[b, :num_pages].long()] + keys = pages[:, : page * head_dim].contiguous().view(torch.float8_e4m3fn) + keys = keys.reshape(num_pages * page, head_dim)[:ctx].float().double() + scales = pages[:, page * head_dim :].contiguous().view(torch.float32) + scales = scales.reshape(num_pages * page)[:ctx].double() + qb = q[b].reshape(next_n * num_heads, head_dim).float().double() + wb = weights[b * next_n : (b + 1) * next_n].reshape(next_n * num_heads, 1).double() + for start in range(0, ctx, 16_384): + stop = min(start + 16_384, ctx) + scores = torch.relu(qb @ keys[start:stop].t()) * wb + summed = scores.reshape(next_n, num_heads, stop - start).sum(dim=1) + out[b * next_n : (b + 1) * next_n, start:stop] = ( + summed * scales[start:stop] + ).float() + return out + + def is_valid( + self, + q: torch.Tensor, + kv_cache: torch.Tensor, + weights: torch.Tensor, + context_lens: torch.Tensor, + block_table: torch.Tensor, + schedule_meta: torch.Tensor, + logits: torch.Tensor, + max_context_len: int, + early_kv: bool, + ) -> None: + assert q.dim() == 4 and q.dtype == torch.float8_e4m3fn, ( + f"q must be [B, next_n, H, D] E4M3, got {tuple(q.shape)} {q.dtype}" + ) + batch, next_n, num_heads, head_dim = q.shape + # The kernel takes Q and the weights through compact layouts; a strided view would be read + # as if it were packed. + assert q.is_contiguous(), "q must be contiguous" + assert head_dim == 128, f"head_dim={head_dim}" + assert weights.shape == (batch * next_n, num_heads) and weights.is_contiguous(), ( + f"weights must be a contiguous [B * next_n, H] = [{batch * next_n}, {num_heads}] " + f"tensor, got {tuple(weights.shape)}" + ) + assert weights.dtype == torch.float32, weights.dtype + # A page is read as [page * D FP8 bytes][page * 4 scale bytes]: any other inner layout + # reads keys and scales from the wrong bytes. + assert kv_cache.dim() == 4 and kv_cache.dtype == torch.uint8, ( + "kv_cache must be the [pages, page, 1, D + 4] uint8 view of the indexer K cache" + ) + page = kv_cache.shape[1] + assert page in (32, 64, 128) and kv_cache.shape[2:] == (1, head_dim + 4), ( + f"kv_cache pages must be [{page}, 1, {head_dim + 4}] with page in (32, 64, 128)" + ) + assert kv_cache[0].is_contiguous(), "the bytes of one page must be contiguous" + assert context_lens.shape == (batch,) and context_lens.dtype == torch.int32, ( + "context_lens must be [B] int32" + ) + assert block_table.dim() == 2 and block_table.shape[0] == batch, ( + "block_table must hold one row per request" + ) + assert block_table.dtype == torch.int32 and block_table.is_contiguous(), ( + "block_table must be contiguous int32 rows" + ) + lens = context_lens.tolist() + longest = max(lens) if lens else 0 + # The kernel reads whole 128-key compute blocks of block-table entries. + needed = -(-longest // 128) * (128 // page) + assert block_table.shape[1] >= needed, ( + f"block_table rows of {block_table.shape[1]} pages cannot cover {longest} keys " + f"({needed} page entries are read)" + ) + if needed: + assert int(block_table[:, :needed].max()) < kv_cache.shape[0], ( + "a block-table entry points past the cache" + ) + assert longest <= max_context_len, ( + f"context length {longest} exceeds max_context_len={max_context_len}" + ) + sms = num_sms(q.device.index if q.device.index is not None else 0) + assert schedule_meta.shape == (sms + 1, 2) and schedule_meta.dtype == torch.int32, ( + f"schedule_meta must be the [{sms + 1}, 2] int32 schedule of this device's {sms} SMs" + ) + assert logits.dim() == 2 and logits.dtype == torch.float32 and logits.stride(1) == 1, ( + "logits must be an fp32 matrix with a unit column stride" + ) + assert logits.shape[0] == batch * next_n, ( + f"logits must hold B * next_n = {batch * next_n} rows, got {logits.shape[0]}" + ) + assert logits.shape[1] >= -(-max_context_len // SPLIT_KV) * SPLIT_KV, ( + f"logits rows of {logits.shape[1]} columns cannot take the whole {SPLIT_KV}-key " + f"splits of max_context_len={max_context_len}" + ) + + +fp8_paged_mqa_logits = GlmFp8PagedMqaLogits() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/paged_mqa_logits_schedule.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/paged_mqa_logits_schedule.py new file mode 100644 index 000000000000..c2eee988a5dc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/paged_mqa_logits_schedule.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-SM work partition of the indexer's paged MQA logits (DeepGEMM's paged metadata).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata + +from ._fp8_paged_mqa_logits import SPLIT_KV +from ._fp8_paged_mqa_logits import num_sms as _device_sms + +#: The indexer K cache's page size, which the schedule call takes as its block size. +_PAGE = 64 + + +def _spec(context_lens) -> dict: + return dict(context_lens=tuple(context_lens), block_kv=_PAGE) + + +class GlmPagedMqaLogitsSchedule(OpWrapper): + """``out`` [num_sms + 1, 2] int32: where each persistent CTA of the paged MQA logits starts. + + Request ``r`` with ``context_lens[r]`` keys is ``ceil(context_lens[r] / 256)`` units of work, + one per 256-key split (two 128-key compute blocks, one per math warp group), numbered request by + request. With ``total`` units, CTA ``s`` (0 <= s <= num_sms) starts at unit + ``w = s * (total // num_sms) + min(s, total % num_sms)``: ``out[s] = (r, w - first unit of r)`` + for the request ``r`` holding unit ``w``, and ``(num_requests, 0)`` past the last unit. CTA + ``s`` covers [out[s], out[s + 1]). + + Each request is one query atom (its whole verify window), so ``context_lens`` is passed as + [num_requests, 1]. Returns ``out``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="c1 key shard: one request, a full 33,152-key shard", spec=_spec([33_152])), + Cell(why="an empty shard: one request of one key", spec=_spec([1])), + Cell(why="one split per unit edge: 256 keys exactly", spec=_spec([256])), + Cell( + why="more splits than SMs: the whole 1M-key context on one rank", + spec=_spec([1_048_596]), + ), + Cell( + why="four requests, one shorter than a split", + spec=_spec([3_000, 70_000, 129, 40_000]), + ), + ) + + note = """ + Metadata consumed: none. The target passes the shard's context length (``hist_shard_meta``'s + ``shard_ctx``) or the generation requests' indexer KV lengths, the page size (64) as + ``block_kv``, the device's SM count (the logits kernel runs one CTA per SM) and a persistent + [num_sms + 1, 2] int32 ``out`` (captured CUDA graphs read it). The schedule kernel waits on + ``griddepcontrol.wait`` before it reads the lengths. + + ``block_kv`` must be 32, 64 or 128; it does not change the 256-key split on this + architecture. + """ + + def __call__( + self, context_lens: torch.Tensor, block_kv: int, num_sms: int, out: torch.Tensor + ) -> torch.Tensor: + out.copy_( + get_paged_mqa_logits_metadata(context_lens.view(-1, 1), block_kv, num_sms), + non_blocking=True, + ) + return out + + def reference( + self, context_lens: torch.Tensor, block_kv: int, num_sms: int, out: torch.Tensor + ) -> torch.Tensor: + """The partition, in Python integers; ``out`` only gives the result's device.""" + work = [-(-ctx // SPLIT_KV) for ctx in context_lens.tolist()] + prefix = [] + total = 0 + for units in work: + total += units + prefix.append(total) + quota, rem = divmod(total, num_sms) + rows = [] + for sm in range(num_sms + 1): + w = sm * quota + min(sm, rem) + request = next((r for r, end in enumerate(prefix) if end > w), len(work)) + if request < len(work): + rows.append((request, w - (prefix[request - 1] if request else 0))) + else: + rows.append((len(work), 0)) + return torch.tensor(rows, dtype=torch.int32, device=out.device) + + def is_valid( + self, context_lens: torch.Tensor, block_kv: int, num_sms: int, out: torch.Tensor + ) -> None: + # Each element is read as one request's length. + assert context_lens.dim() == 1, "context_lens must be [num_requests]" + # The schedule is right for any SM count, but the logits kernel launches one CTA per SM + # and reads entries [s] and [s + 1]: another count leaves keys unscored or reads past the + # schedule. + device = context_lens.device.index if context_lens.device.index is not None else 0 + assert num_sms == _device_sms(device), ( + f"num_sms={num_sms} differs from the device's SM count, the logits kernel's CTA count" + ) + # A wider int type would be cast silently by the copy. + assert out.dtype == torch.int32, f"out must be int32, got {out.dtype}" + + +paged_mqa_logits_schedule = GlmPagedMqaLogitsSchedule() diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_fp8_paged_mqa_logits.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_fp8_paged_mqa_logits.py new file mode 100644 index 000000000000..56757d1a6684 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_fp8_paged_mqa_logits.py @@ -0,0 +1,202 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM fp8_paged_mqa_logits entry. + +The keys live in a real INDEX_KEY pool of 64-token pages, scattered over random physical pages; the +schedule comes from the production schedule op (``paged_mqa_logits_schedule``), as at c1. + +Most cells use integer data: Q and K are small integers (exact in E4M3), the key scales powers of +two and the weights small integers times 2^-4. Every product and partial sum is then an exactly +representable fp32 value, whatever the accumulation order, so the kernel must equal the fp64 +reference bit for bit. One cell uses UE8M0-quantized random keys and real-valued weights, where the +accumulation order is visible and a band applies. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.fp8_paged_mqa_logits import ( # noqa: E501 + fp8_paged_mqa_logits as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.paged_mqa_logits_schedule import ( # noqa: E501 + paged_mqa_logits_schedule, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../..", "."] +import _indexer_inputs as inputs # noqa: E402 +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +assert torch.cuda.is_available(), "fp8_paged_mqa_logits requires a CUDA device" + +_SPLIT = 256 + + +def _sms() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + +def _keys(count: int, data: str, gen: torch.Generator): + """(E4M3 keys [count, 128], fp32 scales [count]).""" + if data == "int": + keys = torch.randint(-3, 4, (count, inputs.HEAD_DIM), generator=gen).float() + scales = torch.pow(2.0, torch.randint(-6, 2, (count,), generator=gen).float()) + return keys.to(torch.float8_e4m3fn), scales + raw = torch.randn(count, inputs.HEAD_DIM, generator=gen) + amax = raw.abs().amax(dim=1).clamp_min(1e-4) + scales = torch.pow(2.0, torch.ceil(torch.log2(amax / 448.0))) + return (raw / scales[:, None]).to(torch.float8_e4m3fn), scales + + +def _build(spec, seed: int) -> dict: + gen = torch.Generator().manual_seed(seed) + lens = list(spec["context_lens"]) + batch, next_n, heads, d = len(lens), spec["next_n"], spec["num_heads"], spec["head_dim"] + page = spec["tokens_per_block"] + max_ctx = spec["max_context_len"] + width = -(-max_ctx // page) + pages_per_request = [-(-n // page) for n in lens] + pool_pages = sum(pages_per_request) + 3 + cache, whole = inputs.k_cache_pool(pool_pages, spec["pool_layout"]) + order = torch.randperm(pool_pages, generator=gen).to(torch.int32) + table = torch.zeros(batch, width, dtype=torch.int32) + used = 0 + flat_pages = cache.reshape(pool_pages, -1) + for b, n in enumerate(lens): + count = pages_per_request[b] + table[b, :count] = order[used : used + count] + used += count + keys, scales = _keys(count * page, spec["data"], gen) + page_bytes = torch.cat( + [ + keys.view(torch.uint8).reshape(count, page * d), + scales.view(torch.uint8).reshape(count, page * 4), + ], + dim=1, + ) + flat_pages[table[b, :count].long().cuda()] = page_bytes.cuda() + if spec["data"] == "int": + q = torch.randint(-3, 4, (batch, next_n, heads, d), generator=gen).float() + weights = torch.randint(-3, 4, (batch * next_n, heads), generator=gen).float() * 2.0**-4 + else: + q = torch.randn(batch, next_n, heads, d, generator=gen) + weights = torch.randn(batch * next_n, heads, generator=gen) * 0.05 + context_lens = torch.tensor(lens, dtype=torch.int32).cuda() + schedule = torch.empty(_sms() + 1, 2, dtype=torch.int32, device="cuda") + paged_mqa_logits_schedule(context_lens, page, _sms(), schedule) + logits = torch.full( + (batch * next_n, -(-max_ctx // _SPLIT) * _SPLIT), float("nan"), device="cuda" + ) + return dict( + q=q.to(torch.float8_e4m3fn).cuda(), + kv_cache=cache, + weights=weights.cuda(), + context_lens=context_lens, + block_table=table.cuda(), + schedule_meta=schedule, + logits=logits, + max_context_len=max_ctx, + early_kv=spec["early_kv"], + ) + + +def _compare(out: torch.Tensor, ref: torch.Tensor, lens, next_n: int, exact: bool) -> None: + """Request by request over its keys; the columns past a request's keys are unspecified.""" + for b, n in enumerate(lens): + got = out[b * next_n : (b + 1) * next_n, :n] + want = ref[b * next_n : (b + 1) * next_n, :n] + if exact: + assert torch.equal(got, want), f"request {b}: logits differ from the exact reference" + else: + # fp32 accumulation in another order: products of E4M3 values are exact, so the + # difference is the rounding of the 128-term dots and of the 32-term weighted sum, + # a few ulps of the partial sums. 64 ulps of each row's largest logit leave an order of + # magnitude of margin; a wrong head, weight or key is off by the logits themselves. + assert_within_ulp(got.contiguous(), want.contiguous(), element_ulp=64.0, rms_ulp=8.0) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + ref = op.reference(**args) + out = op(**args) + torch.cuda.synchronize() + assert out.data_ptr() == args["logits"].data_ptr(), "the op writes the caller's buffer" + assert out.shape == (args["logits"].shape[0], cell.spec["max_context_len"]) + _compare(out, ref, cell.spec["context_lens"], cell.spec["next_n"], cell.spec["data"] == "int") + + +def test_early_kv_kernel_equals_the_default_kernel_bitwise(): + """The N split and the early KV loads move work, not arithmetic: same bits on real values.""" + spec = next(c.spec for c in op.CELLS if c.spec["data"] == "randn") + args = _build(spec, seed=51) + early = op(**dict(args, early_kv=True)).clone() + default = op(**dict(args, early_kv=False)) + torch.cuda.synchronize() + n = spec["context_lens"][0] + assert torch.equal(early[:, :n], default[:, :n]) + + +def test_cuda_graph_replay_reads_new_queries_and_keys(): + """The decode path is captured: a replay follows Q, the weights and the cache in place.""" + spec = op.CELLS[0].spec + args = _build(spec, seed=52) + op(**args) # compile outside capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(**args) + fresh = _build(spec, seed=53) + args["q"].copy_(fresh["q"]) + args["weights"].copy_(fresh["weights"]) + args["kv_cache"].copy_(fresh["kv_cache"]) + args["block_table"].copy_(fresh["block_table"]) + ref = op.reference(**args) + graph.replay() + torch.cuda.synchronize() + _compare(args["logits"], ref, spec["context_lens"], spec["next_n"], exact=True) + + +def test_guard_refuses_a_block_table_too_narrow_for_the_keys(): + args = _build(op.CELLS[2].spec, seed=54) + narrow = dict( + args, + context_lens=torch.tensor([5_000], dtype=torch.int32, device="cuda"), + max_context_len=5_000, + block_table=args["block_table"][:, :10].contiguous(), + ) + with pytest.raises(AssertionError, match="cannot cover"): + with validating(op): + op(**narrow) + + +def test_guard_refuses_a_schedule_of_another_sm_count(): + args = _build(op.CELLS[2].spec, seed=55) + schedule = torch.zeros(_sms(), 2, dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match="schedule_meta"): + with validating(op): + op(**dict(args, schedule_meta=schedule)) + + +def test_guard_refuses_logits_without_room_for_whole_splits(): + args = _build(op.CELLS[2].spec, seed=56) + narrow = torch.empty(args["logits"].shape[0], args["max_context_len"] - 64, device="cuda") + with pytest.raises(AssertionError, match="whole"): + with validating(op): + op(**dict(args, logits=narrow)) + + +def test_guard_refuses_strided_weights(): + args = _build(op.CELLS[2].spec, seed=57) + strided = args["weights"].t().contiguous().t() + assert not strided.is_contiguous() + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op(**dict(args, weights=strided)) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_paged_mqa_logits_schedule.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_paged_mqa_logits_schedule.py new file mode 100644 index 000000000000..cbbc5b6b61ba --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_paged_mqa_logits_schedule.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM paged_mqa_logits_schedule entry: the schedule is integers, compared exactly. + +Every cell runs at the device's SM count, the count the logits kernel launches its CTAs with. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.paged_mqa_logits_schedule import ( # noqa: E501 + paged_mqa_logits_schedule as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +assert torch.cuda.is_available(), "paged_mqa_logits_schedule requires a CUDA device" + + +def _sms() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + +def _build(spec) -> dict: + sms = _sms() + return dict( + context_lens=torch.tensor(spec["context_lens"], dtype=torch.int32, device="cuda"), + block_kv=spec["block_kv"], + num_sms=sms, + out=torch.full((sms + 1, 2), -7, dtype=torch.int32, device="cuda"), + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args = _build(cell.spec) + ref = op.reference(**args) + out = op(**args) + torch.cuda.synchronize() + assert out is args["out"] + assert torch.equal(out, ref), f"schedule differs:\n{out.tolist()}\nexpected\n{ref.tolist()}" + + +def test_the_partition_covers_every_split_once(): + """Independent of the reference: consecutive starts tile the requests' splits.""" + lens = [3_000, 70_000, 129, 40_000] + args = _build(dict(context_lens=lens, block_kv=64)) + out = op(**args).tolist() + splits = [-(-n // 256) for n in lens] + flat = [(r, s) for r, count in enumerate(splits) for s in range(count)] + position = {unit: i for i, unit in enumerate(flat)} + starts = [position.get(tuple(row), len(flat)) for row in out] + assert starts[0] == 0 and starts[-1] == len(flat) + assert all(a <= b for a, b in zip(starts, starts[1:])) + + +def test_guard_refuses_another_sm_count(): + args = _build(op.CELLS[0].spec) + sms = args["num_sms"] - 1 + with pytest.raises(AssertionError, match="SM count"): + with validating(op): + op(**dict(args, num_sms=sms, out=torch.empty(sms + 1, 2, dtype=torch.int32).cuda())) + + +def test_guard_refuses_two_dimensional_lengths(): + args = _build(op.CELLS[4].spec) + with pytest.raises(AssertionError, match="num_requests"): + with validating(op): + op(**dict(args, context_lens=args["context_lens"].view(2, 2))) From acbb0e8dbb951c25e327070155c5471cb724cafc Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:19:51 -0700 Subject: [PATCH 33/53] Add the GLM MTP index-share stash and reused-TopK remap entries With MTP index share, draft step 0's indexer computes the TopK of the verify window and the later draft steps reuse each request's row of its last accepted token. stash_last_accepted_rows copies those rows into the shared TopK buffer with one Triton launch (the rows index_select would copy, out-of-range counts clamped). convert_req_index_to_global maps the reused request-local rows to MLA KV pool rows through the trtllm::convert_req_index_to_global op, once per draft loop. Both entries are certified on sm_107, exactly. Signed-off-by: Dongfeng Yu --- .../ops/dsa/_mtp_topk_stash.py | 64 +++++++++ .../ops/dsa/convert_req_index_to_global.py | 130 ++++++++++++++++++ .../ops/dsa/stash_last_accepted_rows.py | 79 +++++++++++ ...modeling_v2_convert_req_index_to_global.py | 84 +++++++++++ ...st_modeling_v2_stash_last_accepted_rows.py | 84 +++++++++++ 5 files changed, 441 insertions(+) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_mtp_topk_stash.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/convert_req_index_to_global.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/stash_last_accepted_rows.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_convert_req_index_to_global.py create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_stash_last_accepted_rows.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_mtp_topk_stash.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_mtp_topk_stash.py new file mode 100644 index 000000000000..4e41c60fedaa --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_mtp_topk_stash.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MTP index share: each request's last accepted TopK row of draft step 0, in one Triton launch. + +With MTP index share, draft step 0's indexer computes the TopK of the verify step's next_n rows, +and the draft steps after it reuse, per request, the row of its last accepted token. The kernel +copies those rows, the rows ``torch.index_select(gen_topk, 0, index)`` would copy for +``index[r] = r * next_n + clamp(num_accepted[r] - 1, 0, next_n - 1)``. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _stash_last_accepted_rows_kernel( + src_ptr, # [num_requests * next_n, n_cols] int32 TopK rows + src_stride, + num_accepted_ptr, # [num_requests] accepted tokens per request + dst_ptr, # [num_requests, n_cols] int32 + dst_stride, + next_n, + n_cols, + BLOCK: tl.constexpr, +): + req = tl.program_id(0) + chunk = tl.program_id(1) + # The request's last accepted row of its next_n: num_accepted - 1, clamped to [0, next_n - 1]. + accepted = tl.load(num_accepted_ptr + req) + offset = tl.minimum(tl.maximum(accepted - 1, 0), next_n - 1) + row = req.to(tl.int64) * next_n + offset.to(tl.int64) + cols = chunk * BLOCK + tl.arange(0, BLOCK) + mask = cols < n_cols + values = tl.load(src_ptr + row * src_stride + cols, mask=mask) + tl.store(dst_ptr + req.to(tl.int64) * dst_stride + cols, values, mask=mask) + + +def launch_stash_last_accepted_rows( + gen_topk: torch.Tensor, num_accepted: torch.Tensor, dst: torch.Tensor, next_n: int +) -> None: + """dst[r] = gen_topk[r * next_n + clamp(num_accepted[r] - 1, 0, next_n - 1)] for each request. + + Args: + gen_topk: [num_requests * next_n, n_cols] int32 TopK rows, unit column stride. + num_accepted: [num_requests] accepted tokens per request (contiguous). + dst: [num_requests, n_cols] int32, unit column stride. + """ + num_requests, n_cols = dst.shape + if num_requests == 0: + return + block = min(triton.next_power_of_2(n_cols), 2048) + _stash_last_accepted_rows_kernel[(num_requests, triton.cdiv(n_cols, block))]( + gen_topk, + gen_topk.stride(0), + num_accepted, + dst, + dst.stride(0), + next_n, + n_cols, + BLOCK=block, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/convert_req_index_to_global.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/convert_req_index_to_global.py new file mode 100644 index 000000000000..eabfcfedf014 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/convert_req_index_to_global.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Request-local TopK token indices to rows of the MLA KV pool, via the trtllm convert op.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +def _spec(rows, **overrides) -> dict: + spec = dict( + rows=rows, + topk=2048, + tokens_per_block=64, + kv_len=530_000, + table_width=8_288, + pool_pages=9_000, + page_index_scale=79, + layer=78, + invalid=0, + table_layout="rows", + ) + spec.update(overrides) + return spec + + +class GlmConvertReqIndexToGlobal(OpWrapper): + """``out[t, c]`` = the KV-pool row of request-local token ``tok = token_indices[t, c]`` of + request ``req_id[t]``: ``block_table[req, tok // block_size] * stride_factor + layer_id * + block_size + tok % block_size``; -1 where ``tok < 0``, ``tok // block_size`` is past the table + row, or the table entry is negative. Columns ``[0, num_topk_tokens)``; returns a fresh tensor + shaped like ``token_indices``. + + For the MLA pool of layer ``l`` the target passes ``stride_factor = page_index_scale * + tokens_per_block`` and ``layer_id = layer_offset`` from the KV cache manager's + ``get_primary_pool_page_index_params(l)``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="c1 draft step 1: the reused TopK row of one request, MTP layer 78", + spec=_spec(1), + ), + Cell( + why="-1 TopK slots of a short request stay -1", + spec=_spec(1, kv_len=1_500, invalid=548), + ), + Cell( + why="four requests, table rows of the generation requests only", + spec=_spec(4, kv_len=70_000, table_width=1_200, table_layout="gen_rows"), + ), + ) + + note = """ + Metadata consumed: none. At c1 draft step 1 the target passes the request map of the draft + step (``req_idx_per_token`` of the generation tokens minus the context count), the generation + requests' rows of the MLA block table, the stashed TopK rows ([num_requests, 2048] int32), + the page size, 2048, and the MTP layer's stride factor and layer offset. A plain launch. + + The op reads ``token_indices`` and the table through ``.contiguous()`` copies when they are + strided. + """ + + def __call__( + self, + req_id: torch.Tensor, + block_table: torch.Tensor, + token_indices: torch.Tensor, + block_size: int, + num_topk_tokens: int, + stride_factor: int, + layer_id: int, + ) -> torch.Tensor: + return torch.ops.trtllm.convert_req_index_to_global( + req_id, block_table, token_indices, block_size, num_topk_tokens, stride_factor, layer_id + ) + + def reference( + self, + req_id: torch.Tensor, + block_table: torch.Tensor, + token_indices: torch.Tensor, + block_size: int, + num_topk_tokens: int, + stride_factor: int, + layer_id: int, + ) -> torch.Tensor: + """The pool rows in int64 arithmetic, cast back to int32.""" + tok = token_indices[:, :num_topk_tokens].long() + page = tok.clamp_min(0) // block_size + valid_page = page < block_table.shape[1] + req = req_id.long()[:, None].expand_as(tok) + base = block_table[req, page.clamp(max=block_table.shape[1] - 1)].long() + rows = base * stride_factor + tok.clamp_min(0) % block_size + layer_id * block_size + valid = (tok >= 0) & valid_page & (base >= 0) + out = token_indices.clone() + out[:, :num_topk_tokens] = torch.where(valid, rows, torch.full_like(rows, -1)).int() + return out + + def is_valid( + self, + req_id: torch.Tensor, + block_table: torch.Tensor, + token_indices: torch.Tensor, + block_size: int, + num_topk_tokens: int, + stride_factor: int, + layer_id: int, + ) -> None: + # Columns past num_topk_tokens are left unwritten in the fresh output. + assert num_topk_tokens == token_indices.shape[1], ( + f"num_topk_tokens={num_topk_tokens} leaves columns of the " + f"{token_indices.shape[1]}-wide output unwritten" + ) + assert req_id.shape[0] == token_indices.shape[0], "one request index per TopK row" + # A request index past the table reads another allocation's memory. + assert int(req_id.max()) < block_table.shape[0] and int(req_id.min()) >= 0, ( + "a request index past the block table" + ) + # The pool row is computed in int32. + largest = int(block_table.max()) * stride_factor + (layer_id + 1) * block_size + assert largest < 2**31, "pool rows overflow int32" + + +convert_req_index_to_global = GlmConvertReqIndexToGlobal() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/stash_last_accepted_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/stash_last_accepted_rows.py new file mode 100644 index 000000000000..02b6d14a1300 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/stash_last_accepted_rows.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MTP index share: draft step 0's TopK row of each request's last accepted token, in one launch.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._mtp_topk_stash import launch_stash_last_accepted_rows + + +def _spec(requests, accepted, **overrides) -> dict: + spec = dict(requests=requests, next_n=8, n_cols=2048, accepted=tuple(accepted), dst="rows") + spec.update(overrides) + return spec + + +class GlmStashLastAcceptedRows(OpWrapper): + """``dst[r] = gen_topk[r * next_n + clamp(num_accepted[r] - 1, 0, next_n - 1)]`` per request. + + Draft step 0 of the MTP draft loop computes the TopK of the verify step's ``next_n`` rows; the + draft steps after it reuse, per request, the row of its last accepted token. ``num_accepted`` + is the verify step's accepted-token count per request (1 to next_n); out-of-range counts clamp. + Returns ``dst``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="c1: one request, 8 TopK rows, 3 tokens accepted", spec=_spec(1, [3])), + Cell(why="c1: every draft token accepted (the last row)", spec=_spec(1, [8])), + Cell(why="counts past the window clamp: 0 and 9", spec=_spec(2, [0, 9])), + Cell( + why="four requests into rows of a wider shared buffer", + spec=_spec(4, [1, 5, 8, 2], dst="wide"), + ), + ) + + note = """ + Metadata consumed: none. The target passes the draft step-0 TopK output + ([num_requests * 8, 2048] int32), the verify step's accepted counts per request (int32 on the + device) and the rows [0, num_requests) of the persistent shared-TopK buffer. A plain launch + (no PDL): at c1 it runs on its own stream behind an event recorded after the TopK, and the + layer joins it at the end of its attention. + """ + + def __call__( + self, gen_topk: torch.Tensor, num_accepted: torch.Tensor, dst: torch.Tensor, next_n: int + ) -> torch.Tensor: + launch_stash_last_accepted_rows(gen_topk, num_accepted, dst, next_n) + return dst + + def reference( + self, gen_topk: torch.Tensor, num_accepted: torch.Tensor, dst: torch.Tensor, next_n: int + ) -> torch.Tensor: + """The selected rows as a fresh tensor.""" + requests = dst.shape[0] + offsets = (num_accepted.long() - 1).clamp(0, next_n - 1) + index = torch.arange(requests, device=gen_topk.device) * next_n + offsets + return gen_topk[index, : dst.shape[1]].clone() + + def is_valid( + self, gen_topk: torch.Tensor, num_accepted: torch.Tensor, dst: torch.Tensor, next_n: int + ) -> None: + requests, n_cols = dst.shape + # Rows are addressed as request * next_n + offset with unit column strides. + assert gen_topk.shape == (requests * next_n, n_cols), ( + f"gen_topk must be [{requests} * {next_n}, {n_cols}], got {tuple(gen_topk.shape)}" + ) + assert gen_topk.dtype == dst.dtype, (gen_topk.dtype, dst.dtype) + assert gen_topk.stride(1) == 1 and dst.stride(1) == 1, "unit column strides" + assert num_accepted.is_contiguous() and num_accepted.numel() >= requests, ( + "num_accepted must hold one contiguous count per request" + ) + + +stash_last_accepted_rows = GlmStashLastAcceptedRows() diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_convert_req_index_to_global.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_convert_req_index_to_global.py new file mode 100644 index 000000000000..e81383ff5d33 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_convert_req_index_to_global.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM convert_req_index_to_global entry: integer pool rows, compared exactly. + +The block table has more rows than requests (the context rows of the engine's table come first +when ``table_layout`` is ``gen_rows``) and -1 entries past each request's pages; the TopK rows hold +request-local tokens below each request's KV length, with ``invalid`` -1 slots at the end. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.convert_req_index_to_global import ( # noqa: E501 + convert_req_index_to_global as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +assert torch.cuda.is_available(), "convert_req_index_to_global requires a CUDA device" + + +def _build(spec, seed: int) -> dict: + gen = torch.Generator().manual_seed(seed) + rows, topk, tpb = spec["rows"], spec["topk"], spec["tokens_per_block"] + width, pool = spec["table_width"], spec["pool_pages"] + pages = -(-spec["kv_len"] // tpb) + table = torch.full((rows + 2, width), -1, dtype=torch.int32) + for r in range(rows + 2): + table[r, :pages] = torch.randint(0, pool, (pages,), generator=gen, dtype=torch.int32) + if spec["table_layout"] == "gen_rows": + # The engine's table holds the context requests' rows first; the op gets the generation + # requests' rows, a view starting past them. + table = table.cuda()[2:] + else: + table = table[:rows].cuda() + # A request shorter than the TopK fills its first kv_len slots; the rest are -1. + valid = topk - spec["invalid"] + tokens = torch.full((rows, topk), -1, dtype=torch.int32) + for r in range(rows): + tokens[r, :valid] = torch.randperm(spec["kv_len"], generator=gen)[:valid].to(torch.int32) + return dict( + req_id=torch.arange(rows, dtype=torch.int32, device="cuda"), + block_table=table, + token_indices=tokens.cuda(), + block_size=tpb, + num_topk_tokens=topk, + stride_factor=spec["page_index_scale"] * tpb, + layer_id=spec["layer"], + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + ref = op.reference(**args) + out = op(**args) + torch.cuda.synchronize() + assert torch.equal(out, ref) + if cell.spec["invalid"]: + assert (out[:, -cell.spec["invalid"] :] == -1).all() + + +def test_guard_refuses_fewer_topk_columns_than_the_rows_hold(): + args = _build(op.CELLS[0].spec, seed=81) + with pytest.raises(AssertionError, match="unwritten"): + with validating(op): + op(**dict(args, num_topk_tokens=1024)) + + +def test_guard_refuses_a_request_past_the_table(): + args = _build(op.CELLS[0].spec, seed=82) + req = torch.full_like(args["req_id"], args["block_table"].shape[0]) + with pytest.raises(AssertionError, match="past the block table"): + with validating(op): + op(**dict(args, req_id=req)) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_stash_last_accepted_rows.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_stash_last_accepted_rows.py new file mode 100644 index 000000000000..c8593238d1bd --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_stash_last_accepted_rows.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM stash_last_accepted_rows entry: a row gather, compared exactly. + +The destination is either its own [requests, 2048] buffer or rows of a wider shared buffer whose +other rows must stay as they were. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.stash_last_accepted_rows import ( # noqa: E501 + stash_last_accepted_rows as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +assert torch.cuda.is_available(), "stash_last_accepted_rows requires a CUDA device" + + +def _build(spec, seed: int): + gen = torch.Generator().manual_seed(seed) + requests, next_n, n_cols = spec["requests"], spec["next_n"], spec["n_cols"] + gen_topk = torch.randint(-1, 1 << 20, (requests * next_n, n_cols), generator=gen) + shared = torch.full((requests + 3, n_cols), -7, dtype=torch.int32, device="cuda") + dst = shared[:requests] if spec["dst"] == "wide" else torch.full_like(shared[:requests], -7) + args = dict( + gen_topk=gen_topk.to(torch.int32).cuda(), + num_accepted=torch.tensor(spec["accepted"], dtype=torch.int32, device="cuda"), + dst=dst, + next_n=next_n, + ) + return args, shared + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args, shared = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + ref = op.reference(**args) + out = op(**args) + torch.cuda.synchronize() + assert out is args["dst"] + assert torch.equal(out, ref) + if cell.spec["dst"] == "wide": + assert (shared[cell.spec["requests"] :] == -7).all(), "rows past the requests moved" + + +def test_rows_follow_the_accepted_counts_in_place_under_replay(): + """Captured once, replayed with new counts: the counts are read on the device.""" + args, _ = _build(op.CELLS[0].spec, seed=71) + op(**args) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(**args) + for accepted in (1, 6, 8): + args["num_accepted"].fill_(accepted) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(args["dst"][0], args["gen_topk"][accepted - 1]) + + +def test_guard_refuses_a_topk_of_another_row_count(): + args, _ = _build(op.CELLS[0].spec, seed=72) + with pytest.raises(AssertionError, match="gen_topk must be"): + with validating(op): + op(**dict(args, gen_topk=args["gen_topk"][:4])) + + +def test_guard_refuses_a_strided_destination(): + args, _ = _build(op.CELLS[0].spec, seed=73) + wide = torch.zeros(1, 2 * 2048, dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match="unit column strides"): + with validating(op): + op(**dict(args, dst=wide[:, ::2])) From d8fdb584c5589850874d29c0bcf291fb89d81465 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:29:07 -0700 Subject: [PATCH 34/53] Run the GLM-5.2 decode step's DSA indexer on this rank's key shard For one generation request the indexer layers take their indexer K and weight rows from fused-A's last 160 columns (an extended weight that the built-in kv_a_proj and indexer projections view), their indexer Q from wq_b with the q_a norm folded and an RoPE/FP8/scale epilogue, and write the indexer K cache from a fused K tail on a side stream. Each rank scores its 1/16 shard of the request's keys and the ranks agree on the exact top-2048 through a histogram exchange; the shard's length, block table and logits schedule are set once per step. The flash decode maps the selected positions to cache rows itself, and the layers without an indexer reuse the last selection, gathering its rows before their wait. Other batches keep the built-in indexer. Signed-off-by: Dongfeng Yu --- .../glm_5_2_nvfp4__sm_107__tp16/modeling.py | 289 ++++++++++++++++-- 1 file changed, 262 insertions(+), 27 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py index d25e57487abd..7b955ddc3cda 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py @@ -52,11 +52,26 @@ from .ops.comm._mnnvl_workspace import hop_supported, hop_workspace, publish_target from .ops.comm.mnnvl_hop import mnnvl_hop from .ops.comm.mnnvl_moe_finalize_hop import mnnvl_moe_finalize_hop +from .ops.dsa._fp8_paged_mqa_logits import num_sms +from .ops.dsa._hist_topk import ( + MAX_SHARD, + TOPK, + get_hist_topk_workspace, + hist_topk_supported, + shard_capacity, +) +from .ops.dsa.fp8_paged_mqa_logits import fp8_paged_mqa_logits +from .ops.dsa.hist_shard_meta import hist_shard_meta +from .ops.dsa.hist_topk import hist_topk +from .ops.dsa.indexer_k_tail import indexer_k_tail +from .ops.dsa.paged_mqa_logits_schedule import paged_mqa_logits_schedule from .ops.embedding.fused_embedding import fused_embedding from .ops.gemm.dense_mlp_small_m import INTER_PER_RANK, dense_mlp_small_m +from .ops.gemm.fused_a_gemm import INDEXER_OUT_FEATURES as FUSED_A_INDEXER_OUT_FEATURES from .ops.gemm.fused_a_gemm import MAX_TOKENS as FUSED_A_MAX_TOKENS from .ops.gemm.fused_a_gemm import OUT_FEATURES as FUSED_A_OUT_FEATURES from .ops.gemm.fused_a_gemm import fused_a_gemm +from .ops.gemm.indexer_q_tail import indexer_q_tail from .ops.gemm.low_m_gemm import low_m_gemm from .ops.gemm.o_proj_publish import o_proj_publish from .ops.gemm.q_b_fp8_rope import q_b_fp8_rope @@ -112,9 +127,17 @@ class DecodeTarget: the stored kv_a_proj_with_mqa weight. q_b, with k_b absorbed and the q_a norm folded, writes the RoPE'd FP8 query in its epilogue, while on the aux stream one kernel normalizes, ropes and appends the KV half to the FP8 - cache. The built-in DSA indexer selects 2048 rows, the flash decode attends - over them at 4 heads with v_b folded in, and o_proj's epilogue publishes - this rank's shard into the all-reduce workspace. + cache. The flash decode attends over the 2048 selected rows at 4 heads with + v_b folded in, and o_proj's epilogue publishes this rank's shard into the + all-reduce workspace. + - The DSA indexer (on the layers that run one) shares fused-A: its K and + weight rows are fused-A's last 160 columns. Its Q comes from wq_b (q_a norm + folded) with an epilogue that ropes, quantizes and scales; its K tail + normalizes, ropes, quantizes and writes the indexer K cache on a side + stream. Each rank scores its 1/16 shard of the context's keys, and the ranks + agree on the exact top-2048 by exchanging histograms. The layers without an + indexer reuse the last selection, which the flash decode maps to cache rows + itself and starts gathering before its wait. - Every all-reduce is a one-shot MNNVL hop that adds the residual and applies the following RMSNorm: post_attention_layernorm after attention, and after the MLP or MoE the next layer's input_layernorm (the model's final norm @@ -194,9 +217,6 @@ def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: # [h * 576, h * 576 + 512) are k_b_trans[h] @ W_q_b[h, :qk_nope] and the next 64 rows # W_q_b[h, qk_nope:], every column times gamma_q_a; FP64, rounded to BF16 once. self._q_b_absorbed = {} - # What fused-A prefetches into L2 while it waits for its input: the weights of the - # attention GEMMs that follow it. - self._l2_after_fused_a = {} for layer in layers: mla = layer.self_attn qk_head_dim = mla.qk_nope_head_dim + mla.qk_rope_head_dim @@ -204,12 +224,54 @@ def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: nope = torch.bmm(mla.k_b_proj_trans.double(), w[:, : mla.qk_nope_head_dim]) absorbed = torch.cat([nope, w[:, mla.qk_nope_head_dim :]], 1) absorbed = absorbed * mla.q_a_layernorm.weight.double() - q_b = absorbed.to(torch.bfloat16).reshape(-1, mla.q_lora_rank).contiguous() - self._q_b_absorbed[layer.layer_idx] = q_b - prefetch = [q_b, mla.v_b_proj, mla.o_proj.weight] - if mla.mqa.indexer is not None: - prefetch.insert(1, mla.mqa.indexer.wq_b.weight) + self._q_b_absorbed[layer.layer_idx] = ( + absorbed.to(torch.bfloat16).reshape(-1, mla.q_lora_rank).contiguous() + ) + + # The indexer layers' fused-A weight: kv_a_proj_with_mqa's rows, then the indexer's wk and + # weights_proj rows (BF16), in one tensor whose first rows the built-in kv_a_proj and + # the indexer's projections then view, so nothing is stored twice. Their indexer Q + # comes from wq_b with the q_a norm weight folded into its columns. + self._fused_a_indexer_weight = {} + self._wq_b_folded = {} + for layer in layers: + mla = layer.self_attn + indexer = mla.mqa.indexer + if indexer is None: + continue + assert ( + indexer.wk.weight.dtype == torch.bfloat16 + and indexer.weights_proj.weight.dtype == torch.bfloat16 + ), "the indexer's wk and weights_proj must be BF16 (TRTLLM_DSA_INDEXER_BF16=1)" + kv_a = mla.kv_a_proj_with_mqa.weight + wk = indexer.wk.weight + extended = torch.cat([kv_a, wk, indexer.weights_proj.weight]).contiguous() + assert extended.shape[0] == FUSED_A_INDEXER_OUT_FEATURES + mla.kv_a_proj_with_mqa.weight.data = extended[: kv_a.shape[0]] + indexer.wk.weight.data = extended[kv_a.shape[0] : kv_a.shape[0] + wk.shape[0]] + indexer.weights_proj.weight.data = extended[kv_a.shape[0] + wk.shape[0] :] + self._fused_a_indexer_weight[layer.layer_idx] = extended + self._wq_b_folded[layer.layer_idx] = ( + indexer.wq_b.weight.double() * mla.q_a_layernorm.weight.double() + ).to(torch.bfloat16) + # What fused-A prefetches into L2 while it waits for its input: the weights of the + # attention GEMMs that follow it. + self._l2_after_fused_a = {} + for layer in layers: + mla = layer.self_attn + prefetch = [self._q_b_absorbed[layer.layer_idx], mla.v_b_proj, mla.o_proj.weight] + if layer.layer_idx in self._wq_b_folded: + prefetch.insert(1, self._wq_b_folded[layer.layer_idx]) self._l2_after_fused_a[layer.layer_idx] = [t for t in prefetch if t.is_contiguous()] + # The key-sharded top-k exchanges histograms over the TP group's MNNVL workspace + # (created here, on every rank in the same order). + self._hist_workspace = ( + get_hist_topk_workspace(self.mapping) if hist_topk_supported(self.mapping) else None + ) + self._hist_buffers = None + self._k_tail_stream = torch.cuda.Stream() + self._k_tail_fork = torch.cuda.Event() + self._k_tail_done = torch.cuda.Event() device = core.embed_tokens.weight.device # The attention scales the KV-half kernel writes and the flash decode reads. @@ -271,8 +333,9 @@ def forward( layer, position_ids, hidden_states, residual, attn_metadata ) return hidden_states + sharded = self._shard_indexer_keys(attn_metadata) for layer in layers: - attn_out = self._attention(layer, position_ids, hidden_states, attn_metadata) + attn_out = self._attention(layer, position_ids, hidden_states, attn_metadata, sharded) hidden_states, residual = self._post_attention_hop(layer, attn_out, residual) if isinstance(layer.mlp, GatedMLP): mlp = layer.mlp @@ -285,11 +348,60 @@ def forward( hidden_states, residual = self._moe(layer, hidden_states, residual) return hidden_states - def _attention(self, layer, position_ids, hidden_states, attn_metadata) -> torch.Tensor: + def _shard_indexer_keys(self, attn_metadata) -> bool: + """Point this rank's indexer at its 1/16 shard of the request's keys for this step: + the shard's length and block table, and the logits kernel's schedule over it. + + Applies to one generation request whose shard fits the top-k kernel; returns whether it + did, and otherwise every indexer layer scores the whole context with the built-in + indexer.""" + manager = attn_metadata.kv_cache_manager + if self._hist_workspace is None or attn_metadata.num_seqs != 1: + return False + if self._hist_buffers is None: + capacity = shard_capacity(manager.max_seq_len, self.mapping.tp_size) + if capacity > MAX_SHARD: + return False + device = self._unit_scale.device + tokens_per_block = manager.tokens_per_block + rows = FUSED_A_MAX_TOKENS + sms = num_sms(device.index) + self._hist_buffers = dict( + capacity=capacity, + shard_ctx=torch.zeros(1, dtype=torch.int32, device=device), + shard_block_table=torch.zeros( + (1, -(-capacity // tokens_per_block)), dtype=torch.int32, device=device + ), + schedule=torch.zeros((sms + 1, 2), dtype=torch.int32, device=device), + logits=torch.empty( + (rows, -(-capacity // 256) * 256), dtype=torch.float32, device=device + ), + topk=torch.empty((rows, TOPK), dtype=torch.int32, device=device), + k_fp8=torch.empty((rows, 128), dtype=torch.float8_e4m3fn, device=device), + k_scale=torch.empty((rows, 1), dtype=torch.float32, device=device), + sms=sms, + ) + buffers = self._hist_buffers + hist_shard_meta( + attn_metadata.gen_indexer_kv_lens_cuda_runtime, + attn_metadata.indexer_k_cache_block_offsets[0], + buffers["shard_ctx"], + buffers["shard_block_table"], + self.mapping.tp_size, + self.mapping.tp_rank, + manager.tokens_per_block, + ) + paged_mqa_logits_schedule( + buffers["shard_ctx"], manager.tokens_per_block, buffers["sms"], buffers["schedule"] + ) + return True + + def _attention(self, layer, position_ids, hidden_states, attn_metadata, sharded): """MLA with the DSA indexer over the generation rows: the attention output ahead of o_proj, [rows, heads * v_head_dim].""" mla = layer.self_attn mqa = mla.mqa + indexer = mqa.indexer manager = attn_metadata.kv_cache_manager num_tokens = attn_metadata.num_tokens seq_len_q = num_tokens // attn_metadata.num_seqs @@ -297,20 +409,25 @@ def _attention(self, layer, position_ids, hidden_states, attn_metadata) -> torch mqa._ensure_rope_table_size(attn_metadata.max_seq_len) cos_sin = mqa.rotary_cos_sin kv_lens = attn_metadata.kv_lens_cuda_runtime + local_layer = mqa.get_local_layer_idx(attn_metadata) if self._kv_pool is None: self._kv_pool = manager.get_unique_primary_pool() + own_indexer = sharded and indexer is not None - kv_a = fused_a_gemm( - hidden_states, - mla.kv_a_proj_with_mqa.weight, - self._l2_after_fused_a[layer.layer_idx], - )[:num_tokens] + weight = ( + self._fused_a_indexer_weight[layer.layer_idx] + if own_indexer + else mla.kv_a_proj_with_mqa.weight + ) + kv_a = fused_a_gemm(hidden_states, weight, self._l2_after_fused_a[layer.layer_idx]) + kv_a = kv_a[:num_tokens] q_a = kv_a[:, : mla.q_lora_rank] + latent_end = mla.q_lora_rank + head_dim quant_q = hidden_states.new_empty( (num_tokens, mla.num_heads_tp * head_dim), dtype=torch.uint8 ) - def q_half(): + def q_b(): q_b_fp8_rope( q_a, self._q_b_absorbed[layer.layer_idx], @@ -327,7 +444,7 @@ def q_half(): def kv_half(): mla_kv_half_append( - kv_a[:, mla.q_lora_rank :], + kv_a[:, mla.q_lora_rank : latent_end], mla.kv_a_layernorm.weight, mla.kv_a_layernorm.variance_epsilon, cos_sin, @@ -337,7 +454,7 @@ def kv_half(): attn_metadata.kv_cache_block_offsets, manager.kv_cache_pool_pointers, manager.kv_cache_pool_mapping, - mqa.get_local_layer_idx(attn_metadata), + local_layer, manager.tokens_per_block, attn_metadata.max_seq_len, self._bmm1_scale, @@ -345,10 +462,50 @@ def kv_half(): mla.softmax_scale, ) + output = hidden_states.new_empty( + (hidden_states.shape[0], mla.num_heads_tp * mla.v_head_dim) + ) + q_heads = quant_q.view(num_tokens, mla.num_heads_tp, head_dim) + if sharded: + + def select_then_q_b(): + if own_indexer: + self._select_rows( + indexer, kv_a, latent_end, q_a, position_ids, attn_metadata, num_tokens, mla + ) + q_b() + + # The KV half on the aux stream beside the indexer and q_b. + maybe_execute_in_parallel( + select_then_q_b, kv_half, mla.ln_events[0], mla.ln_events[1], mla.aux_stream + ) + page_index_scale, layer_offset = manager.get_primary_pool_page_index_params( + local_layer + ) + sparse_mla_flash_decode( + q_heads, + self._kv_pool, + self._hist_buffers["topk"][:num_tokens], + kv_lens, + self._bmm1_scale, + self._bmm2_scale, + mla.v_b_proj, + output, + seq_len_q, + self.index_topk, + attn_metadata.block_table[: attn_metadata.num_seqs], + manager.tokens_per_block, + page_index_scale, + layer_offset, + # The selection of an earlier layer: its rows can load before the wait. + not own_indexer, + ) + return output + + # The built-in indexer over the whole context, then the selection as cache rows. maybe_execute_in_parallel( - q_half, kv_half, mla.ln_events[0], mla.ln_events[1], mla.aux_stream + q_b, kv_half, mla.ln_events[0], mla.ln_events[1], mla.aux_stream ) - indexer = mqa.indexer indexer_intermediates = None if indexer is not None: qr = mla.q_a_layernorm(q_a) @@ -372,11 +529,8 @@ def kv_half(): ), ), ) - output = hidden_states.new_empty( - (hidden_states.shape[0], mla.num_heads_tp * mla.v_head_dim) - ) sparse_mla_flash_decode( - quant_q.view(num_tokens, mla.num_heads_tp, head_dim), + q_heads, self._kv_pool, indices, kv_lens, @@ -391,6 +545,87 @@ def kv_half(): indexer.maybe_join_prev_topk_copy() return output + def _select_rows(self, indexer, kv_a, latent_end, q_a, position_ids, attn_metadata, rows, mla): + """The indexer of a layer that runs one, on this rank's key shard: writes the exact + top-2048 cache positions of every row into the shared selection buffer.""" + buffers = self._hist_buffers + manager = attn_metadata.kv_cache_manager + rope = indexer.rotary_emb + cos_sin = rope.rotary_cos_sin.view(rope.max_positions, -1) + interleave = not rope.is_neox + positions = position_ids.view(-1)[:rows] + index_k = kv_a[:, latent_end : latent_end + indexer.head_dim] + index_weights = kv_a[:, latent_end + indexer.head_dim :] + k_cache = manager.get_indexer_k_cache_buffers(indexer.layer_idx) + + # The K tail on its own stream, beside the Q tail: k_norm, RoPE, FP8 and the K cache. + self._k_tail_fork.record() + with torch.cuda.stream(self._k_tail_stream): + self._k_tail_stream.wait_event(self._k_tail_fork) + indexer_k_tail( + index_k, + positions, + cos_sin, + indexer.k_norm.weight, + indexer.k_norm.bias, + indexer.k_norm.variance_epsilon, + indexer.rope_dim, + interleave, + indexer.scale_fmt == "ue8m0", + buffers["k_fp8"][:rows], + buffers["k_scale"][:rows], + k_cache, + attn_metadata.slot_mapping_fp8[:rows], + attn_metadata.slot_mapping_scale[:rows], + rows, + True, + True, + False, + ) + self._k_tail_done.record() + + q_fp8 = q_a.new_empty((rows, indexer.n_heads, indexer.head_dim), dtype=torch.float8_e4m3fn) + q_scale = q_a.new_empty((rows, indexer.n_heads, 1), dtype=torch.float32) + weights = q_a.new_empty((rows, indexer.n_heads), dtype=torch.float32) + indexer_q_tail( + q_a, + self._wq_b_folded[indexer.layer_idx], + positions, + cos_sin, + index_weights, + q_fp8, + q_scale, + weights, + mla.q_a_layernorm.variance_epsilon, + indexer.n_heads, + indexer.head_dim, + indexer.rope_dim, + interleave, + indexer.scale_fmt == "ue8m0", + indexer.weight_scale_factor, + ) + torch.cuda.current_stream().wait_event(self._k_tail_done) + logits = fp8_paged_mqa_logits( + q_fp8.view(1, rows, indexer.n_heads, indexer.head_dim), + k_cache, + weights, + buffers["shard_ctx"], + buffers["shard_block_table"], + buffers["schedule"], + buffers["logits"][:rows], + buffers["capacity"], + True, + ) + hist_topk( + logits, + attn_metadata.gen_indexer_kv_lens_cuda_runtime, + buffers["topk"][:rows], + self._hist_workspace.comm_buffer, + self._hist_workspace.state, + rows, + self.mapping.tp_size, + ) + def _post_attention_hop(self, layer, attn_out: torch.Tensor, residual: torch.Tensor): """o_proj published into the all-reduce workspace, then the hop with the residual add and post_attention_layernorm: (normed, residual).""" From bfa7142f56260f866e5a8e0efe043e45ec7e6cbf Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:34:46 -0700 Subject: [PATCH 35/53] Compute the GLM-5.2 verify logits on the split-K LM head and the MNNVL all-gather On a generation step with speculative decoding the target's shell computes the verify rows' logits itself: the LM head on this rank's vocabulary shard as a split-K GEMM, then one MNNVL all-gather that writes every rank's shard as FP32, on a Lamport workspace the decode target creates with the others. Rows the kernels do not serve keep the built-in logits processor. Signed-off-by: Dongfeng Yu --- .../glm_5_2_nvfp4__sm_107__tp16/modeling.py | 82 +++++++++++++++++-- 1 file changed, 76 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py index 7b955ddc3cda..c83e9325ceda 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py @@ -34,7 +34,10 @@ from tensorrt_llm._torch.distributed import AllReduceFusionOp, AllReduceParams from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3ForCausalLM, DeepseekV3Model -from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM +from tensorrt_llm._torch.models.modeling_speculative import ( + SpecDecOneEngineForCausalLM, + _slice_spec_position_ids, +) from tensorrt_llm._torch.modules.gated_mlp import GatedMLP from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel from tensorrt_llm._torch.modules.swiglu import swiglu @@ -49,7 +52,13 @@ from . import weights as _weights from .ops.attention.mla_kv_half_append import mla_kv_half_append from .ops.attention.sparse_mla_flash_decode import sparse_mla_flash_decode -from .ops.comm._mnnvl_workspace import hop_supported, hop_workspace, publish_target +from .ops.comm._mnnvl_workspace import ( + LamportExchangeWorkspace, + hop_supported, + hop_workspace, + publish_target, +) +from .ops.comm.mnnvl_allgather import mnnvl_allgather from .ops.comm.mnnvl_hop import mnnvl_hop from .ops.comm.mnnvl_moe_finalize_hop import mnnvl_moe_finalize_hop from .ops.dsa._fp8_paged_mqa_logits import num_sms @@ -85,6 +94,9 @@ # facts, and the constructor states them again for a direct construction. _COMPUTE_CAPABILITY = (10, 7) _TP = 16 +# One Lamport buffer of the verify logits' all-gather: every rank's BF16 vocabulary shard of +# up to 8 rows, with room to spare. +_LOGITS_EXCHANGE_BYTES = 16 << 20 class PrefillTarget: @@ -150,6 +162,9 @@ class DecodeTarget: own stream) right before FC1, runs FC2 on the K = 128 kernel and leaves the finalize step to the post-MoE hop. + The verify logits come from the LM head on this rank's vocabulary shard + (split-K GEMM) and one MNNVL all-gather that widens them to FP32. + A larger batch runs the built-in modules in the same order. """ @@ -269,6 +284,7 @@ def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: get_hist_topk_workspace(self.mapping) if hist_topk_supported(self.mapping) else None ) self._hist_buffers = None + self._logits_workspace = LamportExchangeWorkspace(self.mapping, _LOGITS_EXCHANGE_BYTES) self._k_tail_stream = torch.cuda.Stream() self._k_tail_fork = torch.cuda.Event() self._k_tail_done = torch.cuda.Event() @@ -348,6 +364,23 @@ def forward( hidden_states, residual = self._moe(layer, hidden_states, residual) return hidden_states + def logits(self, lm_head, hidden_states: torch.Tensor) -> Optional[torch.Tensor]: + """FP32 logits over the whole vocabulary of the rows, or None for a batch the kernels + do not serve: this rank's vocabulary shard from the LM head, then the all-gather.""" + rows = hidden_states.shape[0] + weight = lm_head.weight + if not self._fast(rows) or weight.shape[0] * self.mapping.tp_size != lm_head.num_embeddings: + return None + shard = low_m_gemm( + hidden_states.contiguous(), weight, hidden_states.new_empty(rows, weight.shape[0]) + ) + return mnnvl_allgather( + shard, + self._logits_workspace.comm_buffer, + self._logits_workspace.buffer_flags, + self.mapping.tp_size, + ) + def _shard_indexer_keys(self, attn_metadata) -> bool: """Point this rank's indexer at its 1/16 shard of the request's keys for this step: the shard's length and block table, and the logits kernel's schedule over it. @@ -850,6 +883,11 @@ def _select_target(self, attn_metadata: AttentionMetadata) -> "PrefillTarget | D self.post_load_weights() return self._prefill if attn_metadata.num_contexts > 0 else self._decode + def decode_target(self, attn_metadata: AttentionMetadata) -> Optional[DecodeTarget]: + """The decode target when this step runs it, else None.""" + target = self._select_target(attn_metadata) + return target if target is self._decode else None + def forward( self, attn_metadata: AttentionMetadata, @@ -920,16 +958,48 @@ def forward( resource_manager=None, **kwargs, ) -> torch.Tensor: - return super().forward( - attn_metadata=attn_metadata, + decode = self.model.decode_target(attn_metadata) + if decode is None or self.spec_worker is None: + return super().forward( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + return_context_logits=return_context_logits, + spec_metadata=spec_metadata, + resource_manager=resource_manager, + **kwargs, + ) + # A generation step with speculative decoding: the verify logits on the decode + # target's kernels where they serve the rows, then the spec worker. + hidden_states = self.model( input_ids=input_ids, + attn_metadata=attn_metadata, position_ids=position_ids, inputs_embeds=inputs_embeds, - return_context_logits=return_context_logits, spec_metadata=spec_metadata, - resource_manager=resource_manager, **kwargs, ) + if attn_metadata.padded_num_tokens is not None: + num_tokens = attn_metadata.num_tokens + hidden_states = hidden_states[:num_tokens] + if input_ids is not None: + input_ids = input_ids[:num_tokens] + position_ids = _slice_spec_position_ids(position_ids, num_tokens) + rows = hidden_states[spec_metadata.gather_ids] + logits = decode.logits(self.lm_head, rows) + if logits is None: + logits = self.logits_processor.forward(rows, self.lm_head, attn_metadata, True) + return self.spec_worker( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + logits=logits, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=self.draft_model, + resource_manager=resource_manager, + ) def load_weights(self, weights, *args, **kwargs): _weights.load_weights(self, weights) From 5c403ebb266434611f9be7a81d443317e1499da9 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 16:43:31 -0700 Subject: [PATCH 36/53] Build DeepGEMM's JIT kernels for the 10.0 family target on 10.x devices past 10.3 DeepGEMM's JIT targets the device's own architecture (sm_107f on a 10.7 device), for which the bundled CUTLASS enables neither TMA nor tcgen05, so the first kernel that prefetches a TMA descriptor hits a device-side assert (the DSA indexer's logits in the first context warm-up of GLM-5.2). Such devices now build for sm_100f, which they run and CUTLASS knows, the way DeepGEMM already pins 12.x devices to sm_120f. Signed-off-by: Dongfeng Yu --- 3rdparty/fetch_content.json | 5 ++++- 3rdparty/patches/deepgemm_sm10x_family_target.patch | 13 +++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) create mode 100644 3rdparty/patches/deepgemm_sm10x_family_target.patch diff --git a/3rdparty/fetch_content.json b/3rdparty/fetch_content.json index 24c079b57a5a..b9d3f9c4cb06 100644 --- a/3rdparty/fetch_content.json +++ b/3rdparty/fetch_content.json @@ -46,7 +46,10 @@ "git_tag": "b64107f2b9599ca76445b7f62eedf66bae1d095b", "git_submodules_recurse": true, "source_subdir": "dont-add-this-project-with-add-subdirectory", - "patch_file": "patches/deepgemm_configurable_barrier_timeout.patch" + "patch_file": [ + "patches/deepgemm_configurable_barrier_timeout.patch", + "patches/deepgemm_sm10x_family_target.patch" + ] }, { "name": "eigen", diff --git a/3rdparty/patches/deepgemm_sm10x_family_target.patch b/3rdparty/patches/deepgemm_sm10x_family_target.patch new file mode 100644 index 000000000000..44b9282fe76b --- /dev/null +++ b/3rdparty/patches/deepgemm_sm10x_family_target.patch @@ -0,0 +1,13 @@ +--- a/csrc/runtime/jit.hpp ++++ b/csrc/runtime/jit.hpp +@@ -37,6 +37,10 @@ + // Preserve the existing 12.x family target when DeepJIT derives a different minor target. + if (runtime->device.get_arch_major() == 12 and runtime->device.get_arch_minor() != 0) + runtime->default_compiler_options.arch = "120f"; ++ // A 10.x device past 10.3 (e.g. 10.7) builds for the 10.0 family target, which it runs: ++ // the bundled CUTLASS enables TMA and tcgen05 only for the targets it knows. ++ if (runtime->device.get_arch_major() == 10 and runtime->device.get_arch_minor() > 3) ++ runtime->default_compiler_options.arch = "100f"; + return runtime; + }); + } From b3b4c2d04b94d8d696662b2400dd6e51c5efe240 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 17:13:45 -0700 Subject: [PATCH 37/53] Fix the GLM-5.2 tp16 op tests that the first sm_107 run flagged, and add the missing fakes The MTP step tests that certify their cells import the certifying fixture they use. The draft token table's band allows the two BF16 roundings its rows carry (the normalized row and the product), 2 ulp per element. The fused-A GEMM, router GEMV, MNNVL hops and logits all-gather get the fake implementations every trtllm op registers. Signed-off-by: Dongfeng Yu --- .../_torch/custom_ops/cpp_custom_ops.py | 24 +++++++++++++++++++ .../moe/test_modeling_v2_lowm_bf16_moe.py | 4 ++-- ...test_modeling_v2_draft_glue_rmsnorm_cat.py | 4 ++-- .../test_modeling_v2_draft_token_table.py | 9 +++---- ...t_modeling_v2_gather_last_accepted_rows.py | 4 ++-- .../test_modeling_v2_verify_head_accept.py | 4 ++-- 6 files changed, 37 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 563268303c40..fb54c0741a39 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1028,6 +1028,30 @@ def _(input: torch.Tensor, weight: torch.Tensor): def _(x, gate_up_weight, down_weight, gate_up_tiling=-1, down_tiling=-1): return x.new_empty((x.shape[0], down_weight.shape[0])) + @torch.library.register_fake("trtllm::glm_fused_a_gemm") + def _(hidden_states, weight, l2_prefetch): + return hidden_states.new_empty((hidden_states.shape[0], weight.shape[0])) + + @torch.library.register_fake("trtllm::glm_router_gemm") + def _(hidden_states, weight): + return hidden_states.new_empty((hidden_states.shape[0], weight.shape[0]), + dtype=torch.float32) + + @torch.library.register_fake("trtllm::glm_mnnvl_hop") + def _(input, residual, gamma, epsilon, comm_buffer, buffer_flags, + shard_published): + return [torch.empty_like(residual), torch.empty_like(residual)] + + @torch.library.register_fake("trtllm::glm_mnnvl_moe_finalize_hop") + def _(gemm2_output, expert_weights, expanded_idx_to_permuted_idx, addend, + residual, gamma, epsilon, comm_buffer, buffer_flags): + return [torch.empty_like(residual), torch.empty_like(residual)] + + @torch.library.register_fake("trtllm::glm_mnnvl_allgather") + def _(input, comm_buffer, buffer_flags, group_size): + return input.new_empty((input.shape[0], input.shape[1] * group_size), + dtype=torch.float32) + @torch.library.register_fake("trtllm::fp8_quantize_1x128_packed_ue8m0") def _(input: torch.Tensor, use_r128c4_layout: bool = True) -> tuple[torch.Tensor, torch.Tensor]: diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_lowm_bf16_moe.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_lowm_bf16_moe.py index 8bdec01cf856..d8e7de4347f2 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_lowm_bf16_moe.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_lowm_bf16_moe.py @@ -13,7 +13,7 @@ from tensorrt_llm._utils import get_sm_version __extra_import_path__ = ["../../../.."] -from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 — needs the path declared above pytestmark = pytest.mark.skipif( get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" @@ -72,7 +72,7 @@ def _build(spec, seed): @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) certifying(op) out = op(*args, TOP_K, ROUTED_SCALING_FACTOR) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py index ac731ef79291..ee62c3167fe6 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py @@ -13,7 +13,7 @@ from tensorrt_llm._utils import get_sm_version __extra_import_path__ = ["../../../.."] -from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 — needs the path declared above pytestmark = pytest.mark.skipif( get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" @@ -55,7 +55,7 @@ def _build(spec, seed): @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) certifying(op) out = op(*args, True) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_token_table.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_token_table.py index a92b363d627b..f3faf4ffbda7 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_token_table.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_token_table.py @@ -27,9 +27,10 @@ def test_table_is_the_token_half_of_eh_proj(tp_size, tp_rank) -> None: """Each table row is eh_proj's embedding half of one token of the rank's shard. Against an FP32 RMSNorm and product, rounded where the builder rounds (the normalized row and the product to - BF16), the rows differ only where the order of a 6144-term FP32 sum or the kernel's approximate reciprocal - square root moves a BF16 rounding: one ulp of that element. So 1 ulp against the row scale and 0.25 ulp of - relative RMS. The hidden-half rows are a slice of the weight, compared exactly. + BF16), the rows differ only where the kernel's approximate reciprocal square root moves the normalized row's + BF16 rounding (one ulp of an input, carried through the product) and where the order of a 6144-term FP32 sum + moves the product's (one ulp of the element): 2 ulp per element, 0.25 ulp of relative RMS. The hidden-half + rows are a slice of the weight, compared exactly. """ g = torch.Generator(device="cuda").manual_seed(tp_size * 100 + tp_rank) vocab_shard = 96 @@ -51,7 +52,7 @@ def test_table_is_the_token_half_of_eh_proj(tp_size, tp_rank) -> None: with true_fp32_matmul(): expected = (normed @ weight[:, :HIDDEN].float().t()).to(torch.bfloat16) assert table.shape == (vocab_shard, HIDDEN) and table.dtype == torch.bfloat16 - assert_within_ulp(table, expected, element_ulp=1.0, rms_ulp=0.25) + assert_within_ulp(table, expected, element_ulp=2.0, rms_ulp=0.25) share = HIDDEN // tp_size torch.testing.assert_close( hidden_rows, weight[tp_rank * share : (tp_rank + 1) * share, HIDDEN:], rtol=0, atol=0 diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_gather_last_accepted_rows.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_gather_last_accepted_rows.py index c1e5844278c0..3edc7f4423e5 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_gather_last_accepted_rows.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_gather_last_accepted_rows.py @@ -13,7 +13,7 @@ from tensorrt_llm._utils import get_sm_version __extra_import_path__ = ["../../../.."] -from _validating import validating # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402,F401 — needs the path declared above pytestmark = pytest.mark.skipif( get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" @@ -36,7 +36,7 @@ def _build(spec, seed): @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name """A gather moves bytes: the op and the reference agree bit for bit.""" args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) certifying(op) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_verify_head_accept.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_verify_head_accept.py index 6aa319935f24..9cef188ae5b6 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_verify_head_accept.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_verify_head_accept.py @@ -21,7 +21,7 @@ from tensorrt_llm._utils import get_sm_version __extra_import_path__ = ["../../../.."] -from _validating import validating # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402,F401 — needs the path declared above pytestmark = pytest.mark.skipif( get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" @@ -76,7 +76,7 @@ def _build(spec, seed): @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name """Integer outputs: the op and the ATen chain agree exactly.""" spec = cell.spec target, draft, rows, pool, counter = _build(spec, seed=zlib.crc32(cell.why.encode()) % 2**31) From f34af65d5fbb98fe41ab6c2231ecc18a5fb96057 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 17:32:02 -0700 Subject: [PATCH 38/53] Detach the tensors the GLM split-K GEMMs hand to CuTe DSL DLPack refuses to export a tensor that requires grad, and module parameters such as the LM head's weight may still do so; the split-K GEMMs convert a detached view, which shares the storage. Signed-off-by: Dongfeng Yu --- .../glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk.py index 5613cb522764..8c592c735b35 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk.py @@ -1471,7 +1471,9 @@ def _bmm_publish( def _from_dlpack_dynamic(tensor, leading_dim: int, assumed_align: int = 32): - return from_dlpack(tensor, assumed_align=assumed_align).mark_layout_dynamic( + # Detached: DLPack exports no tensor that requires grad, which module parameters (the LM head + # weight, for one) may still do; the view shares the storage. + return from_dlpack(tensor.detach(), assumed_align=assumed_align).mark_layout_dynamic( leading_dim=leading_dim ) From bc0a93551fde38d4bebb2d8219a8351eed6c8856 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 17:43:41 -0700 Subject: [PATCH 39/53] Check and fold the GLM-5.2 decode target's attention weights per layer DecodeTarget asserts each decoder layer's MLA weights and KV cache layout, and derives the operands its attention kernels read (the absorbed q_b, the fused-A weight with the indexer rows, the folded indexer wq_b, the L2 prefetch list), one layer at a time through _check_attention and _fold_attention_weights; the MoE layers' shared-expert check and the shared expert's two split-K GEMMs around SwiGLU are methods as well, so a layer outside the decoder stack can take the same checks, operands and GEMMs. Signed-off-by: Dongfeng Yu --- .../glm_5_2_nvfp4__sm_107__tp16/modeling.py | 177 ++++++++++-------- 1 file changed, 95 insertions(+), 82 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py index c83e9325ceda..6b3cb7898d3c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py @@ -176,25 +176,8 @@ def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: self.index_topk = config.index_topk layers = core.layers[: core.num_hidden_layers] for layer in layers: - mla = layer.self_attn assert layer.allreduce is not None and layer.next_layer_layernorm is not None - assert not mla.apply_rotary_emb and not mla.mapping.has_cp_helix() - assert tuple(mla.kv_a_proj_with_mqa.weight.shape) == ( - FUSED_A_OUT_FEATURES, - config.hidden_size, - ) and mla.kv_a_proj_with_mqa.weight.dtype == torch.bfloat16, ( - f"layer {layer.layer_idx}: kv_a_proj_with_mqa is not the BF16 [q_a; kv_a] weight" - ) - assert ( - getattr(mla.mqa, "has_fp8_kv_cache", False) - and mla.kv_cache_dtype != "fp8_ds_mla" - and mla.k_b_proj_trans.dtype == torch.bfloat16 - and mla.v_b_proj.dtype == torch.bfloat16 - and mla.v_b_proj.is_contiguous() - and mla.q_b_proj.weight.dtype == torch.bfloat16 - and mla.o_proj.weight.dtype == torch.bfloat16 - ), f"layer {layer.layer_idx}: unexpected MLA weights or KV cache layout" - assert layer.input_layernorm.nvfp4_scale is None, "kv_a_proj is not NVFP4" + self._check_attention(layer) if isinstance(layer.mlp, GatedMLP): assert ( layer.mlp_tp_size == _TP @@ -210,15 +193,7 @@ def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: and layer.fusion_config.POST_MOE_FUSION and layer.mapping.is_multi_node() ), f"layer {layer.layer_idx}: unexpected MoE-layer configuration" - shared = layer.mlp.shared_experts - assert ( - shared is not None - and layer.mlp.shared_output_scale is None - and shared._is_plain_swiglu() - and shared.swiglu_limit is None - and shared.gate_up_proj.weight.dtype == torch.bfloat16 - and shared.down_proj.weight.dtype == torch.bfloat16 - ), f"layer {layer.layer_idx}: the shared expert is not a plain BF16 SwiGLU MLP" + self._check_shared_expert(layer) backend = getattr(layer.mlp.experts, "backend", layer.mlp.experts) assert ( backend.has_nvfp4 @@ -228,56 +203,13 @@ def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: and backend.w3_w1_weight.shape[-1] * 2 == config.hidden_size ), f"layer {layer.layer_idx}: routed experts are not unpadded TRTLLM-Gen NVFP4" - # q_b with k_b absorbed and the q_a norm weight folded, per layer: for head h, rows - # [h * 576, h * 576 + 512) are k_b_trans[h] @ W_q_b[h, :qk_nope] and the next 64 rows - # W_q_b[h, qk_nope:], every column times gamma_q_a; FP64, rounded to BF16 once. + # The attention operands the kernels read, per layer index (_fold_attention_weights). self._q_b_absorbed = {} - for layer in layers: - mla = layer.self_attn - qk_head_dim = mla.qk_nope_head_dim + mla.qk_rope_head_dim - w = mla.q_b_proj.weight.view(mla.num_heads_tp, qk_head_dim, mla.q_lora_rank).double() - nope = torch.bmm(mla.k_b_proj_trans.double(), w[:, : mla.qk_nope_head_dim]) - absorbed = torch.cat([nope, w[:, mla.qk_nope_head_dim :]], 1) - absorbed = absorbed * mla.q_a_layernorm.weight.double() - self._q_b_absorbed[layer.layer_idx] = ( - absorbed.to(torch.bfloat16).reshape(-1, mla.q_lora_rank).contiguous() - ) - - # The indexer layers' fused-A weight: kv_a_proj_with_mqa's rows, then the indexer's wk and - # weights_proj rows (BF16), in one tensor whose first rows the built-in kv_a_proj and - # the indexer's projections then view, so nothing is stored twice. Their indexer Q - # comes from wq_b with the q_a norm weight folded into its columns. self._fused_a_indexer_weight = {} self._wq_b_folded = {} - for layer in layers: - mla = layer.self_attn - indexer = mla.mqa.indexer - if indexer is None: - continue - assert ( - indexer.wk.weight.dtype == torch.bfloat16 - and indexer.weights_proj.weight.dtype == torch.bfloat16 - ), "the indexer's wk and weights_proj must be BF16 (TRTLLM_DSA_INDEXER_BF16=1)" - kv_a = mla.kv_a_proj_with_mqa.weight - wk = indexer.wk.weight - extended = torch.cat([kv_a, wk, indexer.weights_proj.weight]).contiguous() - assert extended.shape[0] == FUSED_A_INDEXER_OUT_FEATURES - mla.kv_a_proj_with_mqa.weight.data = extended[: kv_a.shape[0]] - indexer.wk.weight.data = extended[kv_a.shape[0] : kv_a.shape[0] + wk.shape[0]] - indexer.weights_proj.weight.data = extended[kv_a.shape[0] + wk.shape[0] :] - self._fused_a_indexer_weight[layer.layer_idx] = extended - self._wq_b_folded[layer.layer_idx] = ( - indexer.wq_b.weight.double() * mla.q_a_layernorm.weight.double() - ).to(torch.bfloat16) - # What fused-A prefetches into L2 while it waits for its input: the weights of the - # attention GEMMs that follow it. self._l2_after_fused_a = {} for layer in layers: - mla = layer.self_attn - prefetch = [self._q_b_absorbed[layer.layer_idx], mla.v_b_proj, mla.o_proj.weight] - if layer.layer_idx in self._wq_b_folded: - prefetch.insert(1, self._wq_b_folded[layer.layer_idx]) - self._l2_after_fused_a[layer.layer_idx] = [t for t in prefetch if t.is_contiguous()] + self._fold_attention_weights(layer) # The key-sharded top-k exchanges histograms over the TP group's MNNVL workspace # (created here, on every rank in the same order). self._hist_workspace = ( @@ -311,6 +243,85 @@ def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: self._quant_ready = torch.cuda.Event() self._fast_rows = {} + def _check_attention(self, layer) -> None: + """Assert the layer's MLA weights and KV cache are the layout the kernels take.""" + mla = layer.self_attn + assert not mla.apply_rotary_emb and not mla.mapping.has_cp_helix() + assert tuple(mla.kv_a_proj_with_mqa.weight.shape) == ( + FUSED_A_OUT_FEATURES, + self.hidden_size, + ) and mla.kv_a_proj_with_mqa.weight.dtype == torch.bfloat16, ( + f"layer {layer.layer_idx}: kv_a_proj_with_mqa is not the BF16 [q_a; kv_a] weight" + ) + assert ( + getattr(mla.mqa, "has_fp8_kv_cache", False) + and mla.kv_cache_dtype != "fp8_ds_mla" + and mla.k_b_proj_trans.dtype == torch.bfloat16 + and mla.v_b_proj.dtype == torch.bfloat16 + and mla.v_b_proj.is_contiguous() + and mla.q_b_proj.weight.dtype == torch.bfloat16 + and mla.o_proj.weight.dtype == torch.bfloat16 + ), f"layer {layer.layer_idx}: unexpected MLA weights or KV cache layout" + assert layer.input_layernorm.nvfp4_scale is None, "kv_a_proj is not NVFP4" + + @staticmethod + def _check_shared_expert(layer) -> None: + """Assert the MoE layer's shared expert is the plain BF16 SwiGLU MLP the kernels run.""" + shared = layer.mlp.shared_experts + assert ( + shared is not None + and layer.mlp.shared_output_scale is None + and shared._is_plain_swiglu() + and shared.swiglu_limit is None + and shared.gate_up_proj.weight.dtype == torch.bfloat16 + and shared.down_proj.weight.dtype == torch.bfloat16 + ), f"layer {layer.layer_idx}: the shared expert is not a plain BF16 SwiGLU MLP" + + def _fold_attention_weights(self, layer) -> None: + """The layer's attention operands in the forms the kernels read, by layer index.""" + mla = layer.self_attn + index = layer.layer_idx + # q_b with k_b absorbed and the q_a norm weight folded: for head h, rows + # [h * 576, h * 576 + 512) are k_b_trans[h] @ W_q_b[h, :qk_nope] and the next 64 rows + # W_q_b[h, qk_nope:], every column times gamma_q_a; FP64, rounded to BF16 once. + qk_head_dim = mla.qk_nope_head_dim + mla.qk_rope_head_dim + w = mla.q_b_proj.weight.view(mla.num_heads_tp, qk_head_dim, mla.q_lora_rank).double() + nope = torch.bmm(mla.k_b_proj_trans.double(), w[:, : mla.qk_nope_head_dim]) + absorbed = torch.cat([nope, w[:, mla.qk_nope_head_dim :]], 1) + absorbed = absorbed * mla.q_a_layernorm.weight.double() + self._q_b_absorbed[index] = ( + absorbed.to(torch.bfloat16).reshape(-1, mla.q_lora_rank).contiguous() + ) + + # A layer that runs the indexer gets its fused-A weight: kv_a_proj_with_mqa's rows, + # then the indexer's wk and weights_proj rows (BF16), in one tensor whose first rows + # the built-in kv_a_proj and the indexer's projections then view, so nothing is stored + # twice. Its indexer Q comes from wq_b with the q_a norm weight folded into its columns. + indexer = mla.mqa.indexer + if indexer is not None: + assert ( + indexer.wk.weight.dtype == torch.bfloat16 + and indexer.weights_proj.weight.dtype == torch.bfloat16 + ), "the indexer's wk and weights_proj must be BF16 (TRTLLM_DSA_INDEXER_BF16=1)" + kv_a = mla.kv_a_proj_with_mqa.weight + wk = indexer.wk.weight + extended = torch.cat([kv_a, wk, indexer.weights_proj.weight]).contiguous() + assert extended.shape[0] == FUSED_A_INDEXER_OUT_FEATURES + mla.kv_a_proj_with_mqa.weight.data = extended[: kv_a.shape[0]] + indexer.wk.weight.data = extended[kv_a.shape[0] : kv_a.shape[0] + wk.shape[0]] + indexer.weights_proj.weight.data = extended[kv_a.shape[0] + wk.shape[0] :] + self._fused_a_indexer_weight[index] = extended + self._wq_b_folded[index] = ( + indexer.wq_b.weight.double() * mla.q_a_layernorm.weight.double() + ).to(torch.bfloat16) + + # What fused-A prefetches into L2 while it waits for its input: the weights of the + # attention GEMMs that follow it. + prefetch = [self._q_b_absorbed[index], mla.v_b_proj, mla.o_proj.weight] + if index in self._wq_b_folded: + prefetch.insert(1, self._wq_b_folded[index]) + self._l2_after_fused_a[index] = [t for t in prefetch if t.is_contiguous()] + def _fast(self, rows: int) -> bool: """Whether every kernel of the step serves this many rows (decided once per shape).""" fast = self._fast_rows.get(rows) @@ -744,18 +755,9 @@ def routed(): True, ) - def shared() -> torch.Tensor: - expert = mlp.shared_experts - gate_up = expert.gate_up_proj.weight - h = low_m_gemm(hidden_states, gate_up, hidden_states.new_empty(rows, gate_up.shape[0])) - h = swiglu(h) - return low_m_gemm( - h, expert.down_proj.weight, hidden_states.new_empty(rows, self.hidden_size) - ) - (fc2_output, expert_weights, permuted_rows), shared_output = maybe_execute_in_parallel( routed, - shared, + lambda: self._shared_expert(mlp, hidden_states), mlp.event_dict[EventType.Main], mlp.event_dict[EventType.MoeShared], mlp.aux_stream, @@ -774,6 +776,17 @@ def shared() -> torch.Tensor: buffer_flags, ) + def _shared_expert(self, mlp, hidden_states: torch.Tensor) -> torch.Tensor: + """This rank's partial sum of the shared expert: two split-K GEMMs around SwiGLU.""" + expert = mlp.shared_experts + rows = hidden_states.shape[0] + gate_up = expert.gate_up_proj.weight + h = low_m_gemm(hidden_states, gate_up, hidden_states.new_empty(rows, gate_up.shape[0])) + h = swiglu(h) + return low_m_gemm( + h, expert.down_proj.weight, hidden_states.new_empty(rows, self.hidden_size) + ) + @staticmethod def _builtin_layer(layer, position_ids, hidden_states, residual, attn_metadata): """One decoder layer on the built-in modules, for a batch the kernels do not serve.""" From e79fce21e3ee7d3f24060185bb14296ceaa3a3ea Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 17:44:25 -0700 Subject: [PATCH 40/53] Run the GLM-5.2 MTP draft loop on the decode target's kernels With one-model MTP-Eagle the target builds its own worker, MTPDraftWorker, in place of the built-in MTPEagleWorker, with the same draft KV cache choice and draft model. On a generation step of one request whose draft layer shares the target's KV pool, drafting greedily through the configured draft length, the worker runs the draft loop on DecodeTarget's kernels; every other step runs the built-in loop. Draft step 0 runs on the verify rows: one kernel writes this rank's window of eh_proj's input (embedding gather, enorm, hnorm, concat, column split), eh_proj runs on it as a split-K GEMM, and its hop applies input_layernorm with a zero residual. The MTP layer's attention runs as a verify full-indexer layer on this rank's key shard; on a side stream each request's top-k of its last accepted row is stashed and mapped to KV pool rows for the later steps, which run no indexer and attend to that selection, or, with a context short enough that the built-in metadata skips the indexer, to every position up to their own. After the post-attention hop each request keeps the row of its last accepted token. Every step's MoE runs the router GEMV and the low-M BF16 experts beside the split-K shared expert, and the MoE-finalize hop adds both with the residual and applies shared_head.norm. The draft LM head runs on the vocabulary shard. Before every step but the last, the next step's eh_proj hidden half (hnorm inside its split-K GEMM) starts first on the aux stream, and glm_draft_ehproj_table writes the draft token, the next step's eh_proj output and input_layernorm from a token table, its positions and the sequence and KV lengths; the last step's glm_draft_sample_glue writes its token and the next-new-tokens rows. The worker keeps the built-in loop's MTP index-share flags, draft KV cache context and cleanup, the host side of the switch to one row per request, and the per-step update_for_spec_dec, which rebuilds the DSA decode metadata from the lengths the kernels write. The token table needs the whole eh_proj weight and the low-M MoE its experts in up-then-gate rows, so the target's loader takes both from the checkpoint before the built-in loader shards them (the built-in BF16 MoE may keep its experts in a backend layout). The draft state is built once after the weights load; every rank creates the two draft exchange workspaces there in the same order. Signed-off-by: Dongfeng Yu --- .../glm_5_2_nvfp4__sm_107__tp16/modeling.py | 893 +++++++++++++++++- .../glm_5_2_nvfp4__sm_107__tp16/weights.py | 72 +- 2 files changed, 951 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py index 6b3cb7898d3c..e8b9a5e0b080 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py @@ -9,12 +9,16 @@ are laid out exactly as they were written. Forwards that carry context rows (engine warm-up, an aggregated deployment) run that same decoder stack in ``PrefillTarget``. Generation-only forwards -- every serving step of the -generation server -- run ``DecodeTarget``. +generation server -- run ``DecodeTarget``. With one-model MTP-Eagle, the draft +loop of a one-request generation step runs on ``DecodeTarget``'s kernels as +well (``MTPDraftWorker``). """ from __future__ import annotations -from typing import Any, Literal, Optional +import contextlib +from dataclasses import dataclass +from typing import Any, List, Literal, Optional, Tuple import torch from transformers import PretrainedConfig @@ -31,7 +35,9 @@ forward_generation_sparse_attn, ) from tensorrt_llm._torch.attention.backends.sparse.dsa.params import DSABackendForwardArgs +from tensorrt_llm._torch.attention.backends.sparse.params import MTPIndexShareMetadata from tensorrt_llm._torch.distributed import AllReduceFusionOp, AllReduceParams +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3ForCausalLM, DeepseekV3Model from tensorrt_llm._torch.models.modeling_speculative import ( @@ -46,7 +52,8 @@ get_data_or_none, to_trtllm_gen_act_type, ) -from tensorrt_llm._torch.speculative import SpecMetadata +from tensorrt_llm._torch.speculative import MTPEagleWorker, SpecMetadata +from tensorrt_llm._torch.speculative.eagle3 import _reset_mtp_index_share from tensorrt_llm._torch.utils import EventType from . import weights as _weights @@ -69,13 +76,16 @@ hist_topk_supported, shard_capacity, ) +from .ops.dsa.convert_req_index_to_global import convert_req_index_to_global from .ops.dsa.fp8_paged_mqa_logits import fp8_paged_mqa_logits from .ops.dsa.hist_shard_meta import hist_shard_meta from .ops.dsa.hist_topk import hist_topk from .ops.dsa.indexer_k_tail import indexer_k_tail from .ops.dsa.paged_mqa_logits_schedule import paged_mqa_logits_schedule +from .ops.dsa.stash_last_accepted_rows import stash_last_accepted_rows from .ops.embedding.fused_embedding import fused_embedding from .ops.gemm.dense_mlp_small_m import INTER_PER_RANK, dense_mlp_small_m +from .ops.gemm.draft_hidden_half import draft_hidden_half from .ops.gemm.fused_a_gemm import INDEXER_OUT_FEATURES as FUSED_A_INDEXER_OUT_FEATURES from .ops.gemm.fused_a_gemm import MAX_TOKENS as FUSED_A_MAX_TOKENS from .ops.gemm.fused_a_gemm import OUT_FEATURES as FUSED_A_OUT_FEATURES @@ -84,11 +94,20 @@ from .ops.gemm.low_m_gemm import low_m_gemm from .ops.gemm.o_proj_publish import o_proj_publish from .ops.gemm.q_b_fp8_rope import q_b_fp8_rope +from .ops.gemm.router_gemm import router_gemm +from .ops.moe.lowm_bf16_moe import glm_lowm_bf16_moe from .ops.moe.routed_experts import ( FUSED_ROUTER_MAX_TOKENS, make_fused_router_workspace, routed_experts, ) +from .ops.spec._draft_exchange import HIDDEN_SIZE as DRAFT_HIDDEN_SIZE +from .ops.spec._draft_exchange import DraftExchangeWorkspace, ehproj_table_supports_tp +from .ops.spec._draft_token_table import build_draft_token_table +from .ops.spec.draft_ehproj_table import glm_draft_ehproj_table +from .ops.spec.draft_glue_rmsnorm_cat import glm_draft_glue_rmsnorm_cat +from .ops.spec.draft_sample_glue import glm_draft_sample_glue +from .ops.spec.gather_last_accepted_rows import glm_gather_last_accepted_rows # The deployment this target is written for; route() selects it on the same # facts, and the constructor states them again for a direct construction. @@ -125,6 +144,52 @@ def forward( ) +@dataclass(frozen=True, eq=False) +class _DraftState: + """What the MTP draft steps hold for the engine's life (``DecodeTarget.build_draft``).""" + + # The MTP layer (the built-in module at layer num_hidden_layers) and the LM head it shares. + layer: Any + lm_head: Any + # Detached views of the embedding table and the enorm / hnorm weights for the CuTe kernels. + embed_table: torch.Tensor + enorm_weight: torch.Tensor + hnorm_weight: torch.Tensor + # This rank's columns [lo, hi) of eh_proj's concatenated input (eh_proj is row-parallel), + # and whether draft step 0's glue kernel launches with PDL. + window: Tuple[int, int] + pdl: bool + # What fused-A prefetches into L2 on the draft steps that run no indexer. + l2_after_fused_a: List[torch.Tensor] + # The residual of draft step 0's eh_proj hop: zeros, so its residual output is the sum. + zero_residual: torch.Tensor + # glm_draft_ehproj_table's operands: the token half of eh_proj for the tokens of this rank's + # vocabulary shard, and this rank's rows of eh_proj's hidden half. + token_table: torch.Tensor + hidden_rows: torch.Tensor + # The low-M BF16 MoE's experts and routing constants. + w3_w1: torch.Tensor + w2: torch.Tensor + top_k: int + routed_scaling_factor: float + # The post-MoE hop's routing: one slot of weight 1 per row, holding that row (-1 elsewhere). + expert_weights: torch.Tensor + expert_rows: torch.Tensor + # Draft step 0's top-k of the request's last accepted row, and the request's index. + topk_stash: torch.Tensor + request_rows: torch.Tensor + # The exchange workspaces of glm_draft_ehproj_table and glm_draft_sample_glue. + table_exchange: DraftExchangeWorkspace + sample_exchange: DraftExchangeWorkspace + # The side stream that stashes that top-k and maps it to KV pool rows, its fork and ready + # events, and the events of the hidden half, which runs on the MTP layer's MoE aux stream. + stream: torch.cuda.Stream + selection_fork: torch.cuda.Event + selection_ready: torch.cuda.Event + half_start: torch.cuda.Event + half_done: torch.cuda.Event + + class DecodeTarget: """Generation-only forwards: the serving step of the generation server. @@ -166,6 +231,31 @@ class DecodeTarget: (split-K GEMM) and one MNNVL all-gather that widens them to FP32. A larger batch runs the built-in modules in the same order. + + The MTP draft steps of a one-request step run here too (``draft_first_step``, + ``draft_next_step``, ``draft_head``, ``draft_ehproj_table``, ``draft_sample_glue``), + on state that ``build_draft`` makes once: + + - Draft step 0 runs on the verify rows. One kernel writes this rank's window of + eh_proj's input (embedding gather, enorm, hnorm, concat, column split), eh_proj + runs on it as a split-K GEMM, and its hop applies input_layernorm with a zero + residual. The MTP layer's attention is a verify full-indexer layer's; each + request's top-k of its last accepted row is kept, as KV pool rows of the MTP + layer, for the later steps (on a side stream). After the post-attention hop each + request keeps the row of its last accepted token. + - The later steps take eh_proj's output and its input_layernorm from the previous + step's table kernel. They run no indexer: they attend to draft step 0's + selection, or, with a context short enough that the built-in metadata skips the + indexer, to every position up to their own. + - Every step's MoE runs on that one row: the router GEMV and the low-M BF16 experts, + the shared expert on the aux stream, and the hop that adds both and the residual + and applies shared_head.norm. + - The draft LM head runs on this rank's vocabulary shard. Before every step but the + last, the next step's eh_proj hidden half starts first on the aux stream, and + ``glm_draft_ehproj_table`` exchanges the ranks' argmax and writes the draft + token, the next step's eh_proj output from a token table built at load, its + positions and the sequence and KV lengths. The last step's + ``glm_draft_sample_glue`` writes its token and the next-new-tokens rows. """ def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: @@ -242,6 +332,10 @@ def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: self._quant_fork = torch.cuda.Event() self._quant_ready = torch.cuda.Event() self._fast_rows = {} + # Whether the last verify forward put the indexer layers on this rank's key shard; draft + # step 0 of the same step runs its indexer on that shard. + self._keys_sharded = False + self._draft: Optional[_DraftState] = None def _check_attention(self, layer) -> None: """Assert the layer's MLA weights and KV cache are the layout the kernels take.""" @@ -322,6 +416,119 @@ def _fold_attention_weights(self, layer) -> None: prefetch.insert(1, self._wq_b_folded[index]) self._l2_after_fused_a[index] = [t for t in prefetch if t.is_contiguous()] + @torch.no_grad() + def build_draft(self, layer, lm_head, weights: _weights.DraftWeights) -> None: + """Build what the MTP draft steps hold for the engine's life, once. + + ``layer`` is the MTP layer and ``lm_head`` the vocabulary-parallel LM head it shares; + ``weights`` holds the checkpoint tensors the token table and the low-M MoE take. + Collective: every rank creates the two draft exchange workspaces here, in the same + order. The draft steps after the first reuse draft step 0's selection, as the built-in + indexer does under MTP index share; without that share, or where the kernels do not + serve one row, nothing is built and the built-in draft loop runs. + """ + if self._draft is not None: + return + core = self.core + config = core.model_config.pretrained_config + mapping = self.mapping + tp_size, tp_rank = mapping.tp_size, mapping.tp_rank + hidden = self.hidden_size + mla = layer.self_attn + indexer = mla.mqa.indexer + if ( + indexer is None + or not indexer.mtp_index_share + or not ehproj_table_supports_tp(tp_size) + or not hop_supported(1, hidden, tp_size) + ): + return + self._check_attention(layer) + self._check_shared_expert(layer) + columns = 2 * hidden // tp_size + norms = (layer.enorm, layer.hnorm, layer.input_layernorm, layer.shared_head.norm) + assert ( + hidden == DRAFT_HIDDEN_SIZE + and layer.eh_proj.weight.dtype == torch.bfloat16 + and tuple(layer.eh_proj.weight.shape) == (hidden, columns) + and weights.eh_proj.dtype == torch.bfloat16 + and tuple(weights.eh_proj.shape) == (hidden, 2 * hidden) + and all( + norm.weight.dtype == torch.bfloat16 and tuple(norm.weight.shape) == (hidden,) + for norm in norms + ) + and layer.enorm.variance_epsilon == layer.hnorm.variance_epsilon + ), "the MTP layer's eh_proj and norms are not the BF16 GLM-5.2 layout" + experts = config.n_routed_experts + inter = config.moe_intermediate_size // tp_size + gate = layer.mlp.gate + assert ( + config.n_group == 1 + and gate.weight.dtype == torch.bfloat16 + and tuple(gate.weight.shape) == (experts, hidden) + and tuple(weights.w3_w1.shape) == (experts, 2 * inter, hidden) + and tuple(weights.w2.shape) == (experts, hidden, inter) + ), "the MTP layer's MoE is not the one-group BF16 MoE of the low-M kernels" + embed_table = core.embed_tokens.weight + vocab_shard = lm_head.weight.shape[0] + assert ( + lm_head.weight.dtype == torch.bfloat16 + and vocab_shard * tp_size == lm_head.num_embeddings == embed_table.shape[0] + ), "the draft LM head is not an unpadded BF16 vocabulary shard" + + # The collective part first, before anything that can fail on one rank only. + table_exchange = DraftExchangeWorkspace.for_ehproj_table(mapping) + sample_exchange = DraftExchangeWorkspace.for_sample_glue(mapping) + + self._fold_attention_weights(layer) + device = embed_table.device + enorm_weight = layer.enorm.weight.detach() + token_table, hidden_rows = build_draft_token_table( + weights.eh_proj, + embed_table.detach(), + enorm_weight, + layer.enorm.variance_epsilon, + tp_rank, + tp_size, + vocab_shard, + ) + top_k = config.num_experts_per_tok + expert_weights = torch.zeros((1, top_k), dtype=torch.bfloat16) + expert_weights[0, 0] = 1 + expert_rows = torch.full((1, top_k), -1, dtype=torch.int32) + expert_rows[0, 0] = 0 + l2_after_fused_a = [self._q_b_absorbed[layer.layer_idx], mla.v_b_proj, mla.o_proj.weight] + self._draft = _DraftState( + layer=layer, + lm_head=lm_head, + embed_table=embed_table.detach(), + enorm_weight=enorm_weight, + hnorm_weight=layer.hnorm.weight.detach(), + window=(tp_rank * columns, (tp_rank + 1) * columns), + pdl=get_env_enable_pdl(), + l2_after_fused_a=[t for t in l2_after_fused_a if t.is_contiguous()], + zero_residual=torch.zeros( + (FUSED_A_MAX_TOKENS, hidden), dtype=torch.bfloat16, device=device + ), + token_table=token_table, + hidden_rows=hidden_rows, + w3_w1=weights.w3_w1, + w2=weights.w2, + top_k=top_k, + routed_scaling_factor=config.routed_scaling_factor, + expert_weights=expert_weights.to(device), + expert_rows=expert_rows.to(device), + topk_stash=torch.empty((1, self.index_topk), dtype=torch.int32, device=device), + request_rows=torch.zeros(1, dtype=torch.int32, device=device), + table_exchange=table_exchange, + sample_exchange=sample_exchange, + stream=torch.cuda.Stream(), + selection_fork=torch.cuda.Event(), + selection_ready=torch.cuda.Event(), + half_start=torch.cuda.Event(), + half_done=torch.cuda.Event(), + ) + def _fast(self, rows: int) -> bool: """Whether every kernel of the step serves this many rows (decided once per shape).""" fast = self._fast_rows.get(rows) @@ -355,12 +562,14 @@ def forward( residual = inputs_embeds hidden_states = layers[0].input_layernorm(inputs_embeds) if not self._fast(hidden_states.shape[0]): + self._keys_sharded = False for layer in layers: hidden_states, residual = self._builtin_layer( layer, position_ids, hidden_states, residual, attn_metadata ) return hidden_states sharded = self._shard_indexer_keys(attn_metadata) + self._keys_sharded = sharded for layer in layers: attn_out = self._attention(layer, position_ids, hidden_states, attn_metadata, sharded) hidden_states, residual = self._post_attention_hop(layer, attn_out, residual) @@ -392,6 +601,336 @@ def logits(self, lm_head, hidden_states: torch.Tensor) -> Optional[torch.Tensor] self.mapping.tp_size, ) + def drafts(self, rows: int) -> bool: + """Whether this step's MTP draft steps run here: the draft state exists, this step's + verify forward put the indexer on this rank's key shard (one request; draft step 0 + reuses it) and the kernels serve draft step 0's ``rows``.""" + return self._draft is not None and self._keys_sharded and self._fast(rows) + + def draft_first_step( + self, + attn_metadata, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + batch_indices: torch.Tensor, + num_accepted: torch.Tensor, + draft_len: int, + ): + """MTP draft step 0 over the verify rows, through the MTP layer. + + ``input_ids`` are the tokens the verify step accepted at every row, ``hidden_states`` + the target's output rows, ``num_accepted`` the accepted count of each request + (``batch_indices``). Returns ``(hidden, gather_ids, pool_rows)``: the MTP layer's + output (shared_head.norm applied) at each request's last accepted row; that row's index + among the verify rows; and the selection the later steps attend to, as KV pool rows, + or None where they attend to every position (a short context) or no step follows. + """ + draft = self._draft + layer = draft.layer + rows = hidden_states.shape[0] + col_lo, col_hi = draft.window + eh_proj_input = glm_draft_glue_rmsnorm_cat( + draft.embed_table, + input_ids, + hidden_states, + draft.enorm_weight, + draft.hnorm_weight, + layer.enorm.variance_epsilon, + col_lo, + col_hi, + draft.pdl, + ) + partial = low_m_gemm( + eh_proj_input, layer.eh_proj.weight, eh_proj_input.new_empty((rows, self.hidden_size)) + ) + # A zero residual: the residual output is eh_proj's output, the sum of the partials. + hidden_states, residual = self._hop( + layer.input_layernorm, partial, draft.zero_residual[:rows] + ) + attn_out = self._attention(layer, position_ids, hidden_states, attn_metadata, True) + pool_rows = None + if draft_len > 1 and not attn_metadata.skip_indexer_for_gen_reqs: + pool_rows = self._stash_selection(attn_metadata, num_accepted, draft_len + 1) + hidden_states, residual = self._post_attention_hop(layer, attn_out, residual) + gather_ids, hidden_states, residual = glm_gather_last_accepted_rows( + hidden_states, batch_indices, num_accepted, draft_len + 1, 0, second=residual + ) + return self._draft_moe(hidden_states, residual), gather_ids, pool_rows + + def draft_next_step( + self, + attn_metadata, + eh_proj_output: Tuple[torch.Tensor, torch.Tensor], + pool_rows: Optional[torch.Tensor], + join_selection: bool, + ) -> torch.Tensor: + """An MTP draft step after the first, one row per request, through the MTP layer. + + ``eh_proj_output`` is the previous step's ``(input_layernorm(x), x)`` for this step's + eh_proj output x. The layer runs no indexer. Its rows attend to ``pool_rows``, draft + step 0's selection (``join_selection``, on the first such step, joins the side stream + that made it); with a context short enough that the built-in metadata skips the + indexer, to every position up to their own, the selection that metadata lists (rebuilt + for this step by ``update_for_spec_dec``). Returns the MTP layer's output rows. + """ + draft = self._draft + layer = draft.layer + hidden_states, residual = eh_proj_output + if attn_metadata.skip_indexer_for_gen_reqs: + positions = attn_metadata.topk_indices_buffer[: hidden_states.shape[0]] + attn_out = self._attention( + layer, None, hidden_states, attn_metadata, True, positions=positions + ) + else: + if join_selection: + torch.cuda.current_stream().wait_event(draft.selection_ready) + attn_out = self._attention( + layer, None, hidden_states, attn_metadata, True, pool_rows=pool_rows + ) + hidden_states, residual = self._post_attention_hop(layer, attn_out, residual) + return self._draft_moe(hidden_states, residual) + + def draft_head(self, hidden_states: torch.Tensor, next_step: bool): + """The draft LM head on this rank's vocabulary shard: ``(logits, hidden_half)``. + + With ``next_step`` the next step's eh_proj hidden half is enqueued first, on the aux + stream: this rank's FP32 share of ``W_eh[:, H:] @ hnorm(hidden_states)``, hnorm inside + the split-K GEMM (``draft_ehproj_table`` joins it); else ``hidden_half`` is None. + """ + draft = self._draft + layer = draft.layer + rows = hidden_states.shape[0] + hidden_half = None + if next_step: + hidden_half = torch.empty( + (rows, draft.hidden_rows.shape[0]), + dtype=torch.float32, + device=hidden_states.device, + ) + aux_stream = layer.mlp.aux_stream + draft.half_start.record() + with torch.cuda.stream(aux_stream): + aux_stream.wait_event(draft.half_start) + draft_hidden_half( + hidden_states, + draft.hidden_rows, + hidden_half, + draft.hnorm_weight, + layer.hnorm.variance_epsilon, + ) + draft.half_done.record() + weight = draft.lm_head.weight + logits = low_m_gemm(hidden_states, weight, hidden_states.new_empty(rows, weight.shape[0])) + return logits, hidden_half + + def draft_ehproj_table( + self, + logits: torch.Tensor, + hidden_half: torch.Tensor, + draft_tokens: torch.Tensor, + step: int, + positions: torch.Tensor, + gather_ids: Optional[torch.Tensor], + seq_lens: torch.Tensor, + kv_lens: torch.Tensor, + rewind_accepted: Optional[torch.Tensor], + ): + """The end of draft step ``step`` when another follows, in one kernel per rank. + + The ranks exchange their argmax with its token-table row and their share of the hidden + half (joined here). Every rank writes the draft token into ``draft_tokens[:, step]``, + the next step's eh_proj output x and ``input_layernorm(x)``, and its positions: + ``positions`` (at ``gather_ids`` on draft step 0) plus one. On draft step 0 + ``seq_lens`` becomes 1 and ``kv_lens`` drops the rejected drafts + (``rewind_accepted``); afterwards ``kv_lens`` grows by one. The DSA decode metadata + is left to ``update_for_spec_dec``. Returns ``((normed, x), positions)`` for the next + step. + """ + draft = self._draft + norm = draft.layer.input_layernorm + rows = logits.shape[0] + normed = logits.new_empty((rows, self.hidden_size)) + residual = logits.new_empty((rows, self.hidden_size)) + next_positions = positions.new_empty((rows,)) + first = step == 0 + torch.cuda.current_stream().wait_event(draft.half_done) + glm_draft_ehproj_table( + logits, + draft.table_exchange.buffer, + draft.table_exchange.flags, + self.mapping.tp_size, + draft_tokens, + step, + hidden_half, + draft.token_table, + norm.weight, + norm.variance_epsilon, + normed, + residual, + positions, + next_positions, + gather_ids, + seq_lens, + first, + kv_lens, + 2 if first else 1, + rewind_accepted, + draft_tokens.shape[1], + None, + [], + None, + None, + None, + None, + None, + ) + return (normed, residual), next_positions + + def draft_sample_glue( + self, + logits: torch.Tensor, + draft_tokens: torch.Tensor, + step: int, + seq_lens: torch.Tensor, + kv_lens: torch.Tensor, + rewind_accepted: Optional[torch.Tensor], + accepted_tokens: torch.Tensor, + batch_indices: torch.Tensor, + num_accepted: torch.Tensor, + ) -> torch.Tensor: + """The end of the last draft step ``step``, in one kernel per rank. + + The ranks exchange their argmax; every rank writes the draft token into + ``draft_tokens[:, step]``, the lengths as ``draft_ehproj_table`` does, and returns the + next-new-tokens rows: each request's last accepted token, then its draft tokens. + """ + draft = self._draft + rows = logits.shape[0] + draft_len = draft_tokens.shape[1] + next_new_tokens = draft_tokens.new_empty((rows, draft_len + 1)) + first = step == 0 + glm_draft_sample_glue( + logits, + draft.sample_exchange.buffer, + draft.sample_exchange.flags, + self.mapping.tp_size, + draft_tokens, + step, + # No eh_proj input window and no positions: no step follows. + None, + None, + None, + None, + 0.0, + 0, + 0, + None, + None, + None, + None, + seq_lens, + first, + kv_lens, + 2 if first else 1, + rewind_accepted, + draft_len, + None, + [], + None, + None, + None, + None, + None, + accepted_tokens, + batch_indices, + num_accepted, + next_new_tokens, + ) + return next_new_tokens + + def _stash_selection(self, attn_metadata, num_accepted: torch.Tensor, next_n: int): + """The selection the MTP draft steps after the first attend to, made on the side + stream after draft step 0's attention: each request's top-k of its last accepted row + (the MTP index share), as KV pool rows of the MTP layer. The first later step joins + it (``selection_ready``).""" + draft = self._draft + manager = attn_metadata.kv_cache_manager + requests = num_accepted.shape[0] + tokens_per_block = manager.tokens_per_block + local_layer = draft.layer.self_attn.mqa.get_local_layer_idx(attn_metadata) + page_index_scale, layer_offset = manager.get_primary_pool_page_index_params(local_layer) + main_stream = torch.cuda.current_stream() + draft.selection_fork.record() + with torch.cuda.stream(draft.stream): + draft.stream.wait_event(draft.selection_fork) + stashed = stash_last_accepted_rows( + self._hist_buffers["topk"][: requests * next_n], + num_accepted, + draft.topk_stash[:requests], + next_n, + ) + pool_rows = convert_req_index_to_global( + draft.request_rows[:requests], + attn_metadata.block_table[:requests], + stashed, + tokens_per_block, + self.index_topk, + page_index_scale * tokens_per_block, + layer_offset, + ) + draft.selection_ready.record() + # Allocated on the side stream and read on this one: outside graph capture the memory + # must outlive this stream's reads (a captured graph keeps it). + if not torch.cuda.is_current_stream_capturing(): + pool_rows.record_stream(main_stream) + return pool_rows + + def _draft_moe(self, hidden_states: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: + """The MTP layer's BF16 MoE on one row and the post-MoE hop, which adds the routed and + shared outputs (the routed row as the one slot of weight 1) and the residual and + applies shared_head.norm: the normalized row.""" + draft = self._draft + layer = draft.layer + mlp = layer.mlp + gate = mlp.gate + + def routed() -> torch.Tensor: + router_logits = router_gemm(hidden_states, gate.weight) + return glm_lowm_bf16_moe( + hidden_states, + router_logits, + gate.e_score_correction_bias, + draft.w3_w1, + draft.w2, + draft.top_k, + draft.routed_scaling_factor, + ) + + routed_output, shared_output = maybe_execute_in_parallel( + routed, + lambda: self._shared_expert(mlp, hidden_states), + mlp.event_dict[EventType.Main], + mlp.event_dict[EventType.MoeShared], + mlp.aux_stream, + ) + norm = layer.shared_head.norm + comm_buffer, buffer_flags = hop_workspace( + self.mapping, hidden_states.shape[0], self.hidden_size + ) + normed, _ = mnnvl_moe_finalize_hop( + routed_output, + draft.expert_weights, + draft.expert_rows, + shared_output, + residual, + norm.weight, + norm.variance_epsilon, + comm_buffer, + buffer_flags, + ) + return normed + def _shard_indexer_keys(self, attn_metadata) -> bool: """Point this rank's indexer at its 1/16 shard of the request's keys for this step: the shard's length and block table, and the logits kernel's schedule over it. @@ -440,9 +979,22 @@ def _shard_indexer_keys(self, attn_metadata) -> bool: ) return True - def _attention(self, layer, position_ids, hidden_states, attn_metadata, sharded): + def _attention( + self, + layer, + position_ids, + hidden_states, + attn_metadata, + sharded, + pool_rows: Optional[torch.Tensor] = None, + positions: Optional[torch.Tensor] = None, + ): """MLA with the DSA indexer over the generation rows: the attention output ahead of - o_proj, [rows, heads * v_head_dim].""" + o_proj, [rows, heads * v_head_dim]. + + On an MTP draft step after the first (``sharded``) the rows attend to a given + selection and the indexer does not run: ``pool_rows``, KV pool rows, or + ``positions``, sequence positions the flash decode maps to pool rows itself.""" mla = layer.self_attn mqa = mla.mqa indexer = mqa.indexer @@ -456,14 +1008,20 @@ def _attention(self, layer, position_ids, hidden_states, attn_metadata, sharded) local_layer = mqa.get_local_layer_idx(attn_metadata) if self._kv_pool is None: self._kv_pool = manager.get_unique_primary_pool() - own_indexer = sharded and indexer is not None + reuses_selection = pool_rows is not None or positions is not None + own_indexer = sharded and indexer is not None and not reuses_selection weight = ( self._fused_a_indexer_weight[layer.layer_idx] if own_indexer else mla.kv_a_proj_with_mqa.weight ) - kv_a = fused_a_gemm(hidden_states, weight, self._l2_after_fused_a[layer.layer_idx]) + l2_prefetch = ( + self._draft.l2_after_fused_a + if reuses_selection + else self._l2_after_fused_a[layer.layer_idx] + ) + kv_a = fused_a_gemm(hidden_states, weight, l2_prefetch) kv_a = kv_a[:num_tokens] q_a = kv_a[:, : mla.q_lora_rank] latent_end = mla.q_lora_rank + head_dim @@ -523,13 +1081,27 @@ def select_then_q_b(): maybe_execute_in_parallel( select_then_q_b, kv_half, mla.ln_events[0], mla.ln_events[1], mla.aux_stream ) + if pool_rows is not None: + sparse_mla_flash_decode( + q_heads, + self._kv_pool, + pool_rows, + kv_lens, + self._bmm1_scale, + self._bmm2_scale, + mla.v_b_proj, + output, + seq_len_q, + self.index_topk, + ) + return output page_index_scale, layer_offset = manager.get_primary_pool_page_index_params( local_layer ) sparse_mla_flash_decode( q_heads, self._kv_pool, - self._hist_buffers["topk"][:num_tokens], + self._hist_buffers["topk"][:num_tokens] if positions is None else positions, kv_lens, self._bmm1_scale, self._bmm2_scale, @@ -541,8 +1113,8 @@ def select_then_q_b(): manager.tokens_per_block, page_index_scale, layer_offset, - # The selection of an earlier layer: its rows can load before the wait. - not own_indexer, + # The selection of an earlier indexer layer: its rows can load before the wait. + positions is None and not own_indexer, ) return output @@ -901,6 +1473,12 @@ def decode_target(self, attn_metadata: AttentionMetadata) -> Optional[DecodeTarg target = self._select_target(attn_metadata) return target if target is self._decode else None + def build_draft(self, layer, lm_head, weights: _weights.DraftWeights) -> None: + """Build the decode target's MTP draft-step state (``DecodeTarget.build_draft``).""" + if self._decode is None: + self.post_load_weights() + self._decode.build_draft(layer, lm_head, weights) + def forward( self, attn_metadata: AttentionMetadata, @@ -920,6 +1498,277 @@ def forward( ) +class MTPDraftWorker(MTPEagleWorker): + """The one-model MTP-Eagle worker: the draft loop of a one-request generation step runs + on the decode target's kernels (``DecodeTarget``'s draft steps); every other step runs the + built-in loop. + + The loop keeps the built-in loop's attention-metadata handling: the MTP index-share flags, + the switch to one row per request after draft step 0, the per-step + ``update_for_spec_dec`` that rebuilds the DSA decode metadata from the sequence and KV + lengths the draft-step kernels write, and the cleanup. The draft tokens and the + next-new-tokens rows come from those kernels. + """ + + def __init__( + self, + spec_config, + model_config: ModelConfig, + use_separate_draft_kv_cache: bool, + *, + mapping, + decode_target, + ) -> None: + super().__init__( + spec_config, model_config, use_separate_draft_kv_cache, mapping=mapping + ) + # The core's per-step lookup of the decode target (a bound method rather than the + # core itself, which belongs to the shell and is no submodule of this worker). + self._decode_target = decode_target + # The draft tokens of the last draft loop on the decode target and the next-new-tokens + # rows its last step wrote, for _prepare_next_new_tokens. + self._drafted_rows: Optional[Tuple[torch.Tensor, torch.Tensor]] = None + + def prepare_1st_drafter_inputs( + self, + input_ids, + position_ids, + hidden_states, + accepted_tokens, + attn_metadata, + spec_metadata, + draft_model, + ): + """Without context requests draft step 0's input ids are the accepted-token rows as + they are, flattened, rather than a copy behind an empty context part.""" + if not self.is_mtp_eagle or attn_metadata.num_contexts > 0: + return super().prepare_1st_drafter_inputs( + input_ids, + position_ids, + hidden_states, + accepted_tokens, + attn_metadata, + spec_metadata, + draft_model, + ) + return { + "input_ids": accepted_tokens[:, : spec_metadata.runtime_draft_len + 1].flatten(), + "position_ids": position_ids, + "hidden_states": hidden_states, + "attn_metadata": attn_metadata, + "spec_metadata": spec_metadata, + } + + def _prepare_next_new_tokens( + self, + accepted_tokens, + next_draft_tokens, + batch_indices_cuda, + batch_size, + num_accepted_tokens, + ): + """The rows the last draft-step kernel wrote when these draft tokens are that loop's.""" + drafted, self._drafted_rows = self._drafted_rows, None + if drafted is not None and drafted[0] is next_draft_tokens: + return drafted[1] + return super()._prepare_next_new_tokens( + accepted_tokens, + next_draft_tokens, + batch_indices_cuda, + batch_size, + num_accepted_tokens, + ) + + def _forward_linear_draft_loop( + self, + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + ): + self._drafted_rows = None + decode = self._drafting_target( + inputs, + attn_metadata, + spec_metadata, + draft_kv_cache_manager, + num_contexts, + batch_size, + num_accepted_tokens, + ) + if decode is None: + return super()._forward_linear_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + ) + return self._draft_on_target( + decode, + inputs, + attn_metadata, + spec_metadata, + draft_kv_cache_manager, + batch_size, + num_accepted_tokens, + ) + + def _drafting_target( + self, + inputs, + attn_metadata, + spec_metadata, + draft_kv_cache_manager, + num_contexts, + batch_size, + num_accepted_tokens, + ) -> Optional[DecodeTarget]: + """The decode target when its kernels draft this step, else None. + + They draft a generation step of one request whose draft layer shares the target's KV + pool, through the configured draft length with no draft-length schedule (the draft + steps after the first skip the MTP layer's indexer K write: the next step's draft + step 0 writes those positions before it reads them), greedily (argmax drafts), with + nothing that rewrites the drafts (guided decoding, SA, a draft vocabulary map), and + when the decode target's kernels serve the rows. + """ + draft_len = spec_metadata.runtime_draft_len + rows = batch_size * (draft_len + 1) + input_ids = inputs["input_ids"] + position_ids = inputs["position_ids"] + if not ( + self.is_mtp_eagle + and num_contexts == 0 + and batch_size == 1 + and draft_kv_cache_manager is None + and draft_len == self.max_draft_len + and getattr(self.spec_config, "draft_len_schedule", None) is None + and self.guided_decoder is None + and self.sa_enhancer is None + and self._d2t is None + and not spec_metadata.wants_advanced_draft_sampling + and not spec_metadata.is_ragged_verify + and attn_metadata.num_tokens == rows + and inputs["hidden_states"].shape[0] == rows + and input_ids.dtype == torch.int32 + and tuple(input_ids.shape) == (rows,) + and input_ids.is_contiguous() + and position_ids.dtype in (torch.int32, torch.int64) + and tuple(position_ids.shape) == (rows,) + and position_ids.is_contiguous() + and num_accepted_tokens.dtype == torch.int32 + and num_accepted_tokens.is_contiguous() + and spec_metadata.batch_indices_cuda.dtype == torch.int32 + ): + return None + decode = self._decode_target(attn_metadata) + if decode is None or not decode.drafts(rows): + return None + return decode + + def _draft_on_target( + self, + decode: DecodeTarget, + inputs, + attn_metadata, + spec_metadata, + draft_kv_cache_manager, + batch_size, + num_accepted_tokens, + ) -> torch.Tensor: + """The draft loop on the decode target's kernels: the draft tokens [batch, draft_len]. + + After each step's kernel, which writes the next step's sequence and KV lengths, + ``update_for_spec_dec`` rebuilds the metadata derived from them, as in the built-in + loop; the host side of draft step 0's switch to one row per request comes before it. + """ + draft_len = spec_metadata.runtime_draft_len + # The accepted tokens' rows: draft step 0's input ids (prepare_1st_drafter_inputs). + accepted_tokens = inputs["input_ids"].view(batch_size, draft_len + 1) + batch_indices = spec_metadata.batch_indices_cuda[:batch_size] + num_accepted = num_accepted_tokens[:batch_size] + draft_tokens = accepted_tokens.new_empty((batch_size, draft_len)) + positions = inputs["position_ids"] + eh_proj_output = gather_ids = pool_rows = next_new_tokens = None + uses_mtp_index_share = isinstance(attn_metadata, MTPIndexShareMetadata) + with contextlib.ExitStack() as draft_scope: + if uses_mtp_index_share: + attn_metadata.set_in_mtp_draft_loop(True) + attn_metadata.set_mtp_num_accepted(num_accepted_tokens) + draft_scope.callback(_reset_mtp_index_share, attn_metadata) + attn_metadata = draft_scope.enter_context( + self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager) + ) + # The draft loop's copy of the sequence lengths (prepare_for_spec_dec) and the KV + # lengths, which the draft-step kernels advance. + seq_lens = attn_metadata.seq_lens_cuda[:batch_size] + kv_lens = attn_metadata.kv_lens_cuda[:batch_size] + for i in range(draft_len): + if uses_mtp_index_share: + attn_metadata.set_skip_topk(i > 0) + if i == 0: + hidden_states, gather_ids, pool_rows = decode.draft_first_step( + attn_metadata, + inputs["input_ids"], + positions, + inputs["hidden_states"], + batch_indices, + num_accepted, + draft_len, + ) + else: + hidden_states = decode.draft_next_step( + attn_metadata, eh_proj_output, pool_rows, join_selection=i == 1 + ) + last = i == draft_len - 1 + logits, hidden_half = decode.draft_head(hidden_states, next_step=not last) + rewind_accepted = num_accepted if i == 0 else None + if i == 0: + # One row per request from here on, on the host (no context requests to + # turn into generation ones); the kernel below fills seq_lens. + attn_metadata._seq_lens[:batch_size].fill_(1) + attn_metadata.on_update() + if last: + next_new_tokens = decode.draft_sample_glue( + logits, + draft_tokens, + i, + seq_lens, + kv_lens, + rewind_accepted, + accepted_tokens, + batch_indices, + num_accepted, + ) + else: + eh_proj_output, positions = decode.draft_ehproj_table( + logits, + hidden_half, + draft_tokens, + i, + positions, + gather_ids if i == 0 else None, + seq_lens, + kv_lens, + rewind_accepted, + ) + attn_metadata.update_for_spec_dec() + if i == 0: + attn_metadata.use_spec_decoding = False + self._drafted_rows = (draft_tokens, next_new_tokens) + return draft_tokens + + class ModelingV2Glm52Nvfp4Sm107Tp16( SpecDecOneEngineForCausalLM[Glm52Nvfp4Sm107Tp16Core, PretrainedConfig] ): @@ -959,6 +1808,23 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: ): self.num_hidden_layers = self.config.num_hidden_layers self.model.layers.extend(self.draft_model.mtp_layers) + if type(self.spec_worker) is MTPEagleWorker: + # The target's worker in place of the built-in one, with the same draft KV cache + # choice and draft model. + worker = MTPDraftWorker( + self.spec_config, + model_config, + self.use_separate_draft_kv_cache, + mapping=model_config.mapping, + decode_target=self.model.decode_target, + ) + worker.set_draft_model(self.draft_model) + builtin = self.spec_worker + self.epilogue[:] = [worker if m is builtin else m for m in self.epilogue] + self.spec_worker = worker + # The MTP layer's checkpoint tensors the draft steps derive operands from, held from + # load_weights to post_load_weights. + self._draft_weights: Optional[_weights.DraftWeights] = None def forward( self, @@ -1015,7 +1881,9 @@ def forward( ) def load_weights(self, weights, *args, **kwargs): - _weights.load_weights(self, weights) + self._draft_weights = _weights.load_weights( + self, weights, isinstance(self.spec_worker, MTPDraftWorker) + ) def setup_aliases(self) -> None: DeepseekV3ForCausalLM.setup_aliases(self) @@ -1023,3 +1891,6 @@ def setup_aliases(self) -> None: def post_load_weights(self): super().post_load_weights() self.model.post_load_weights() + draft_weights, self._draft_weights = self._draft_weights, None + if draft_weights is not None: + self.model.build_draft(self.draft_model.mtp_layers[0], self.lm_head, draft_weights) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py index 1e6491eff23f..5b4913b55107 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py @@ -4,17 +4,23 @@ The target's modules are the built-in GLM-5.2 modules (see ``modeling.py``), so the checkpoint loads through the built-in loader into exactly the shapes -and layouts the context server uses. +and layouts the context server uses. The MTP draft steps also derive operands +from two of the MTP layer's checkpoint tensors (``DraftWeights``), taken before +the built-in loader shards them. """ from __future__ import annotations import copy +from dataclasses import dataclass +from typing import Optional +import torch from transformers import PretrainedConfig from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3WeightLoader +from tensorrt_llm._torch.modules.linear import TensorParallelMode, load_weight_shard def normalize_model_config(model_config: ModelConfig[PretrainedConfig]) -> ModelConfig[PretrainedConfig]: @@ -43,6 +49,66 @@ def normalize_model_config(model_config: ModelConfig[PretrainedConfig]) -> Model return model_config -def load_weights(model, weights) -> None: - """Load the checkpoint into the target's modules with the built-in loader.""" +@dataclass(frozen=True, eq=False) +class DraftWeights: + """The MTP layer's checkpoint tensors the decode target derives draft-step operands from. + + Attributes: + eh_proj: The whole ``[H, 2H]`` eh_proj weight, on the model's device. The built-in + eh_proj is row-parallel (``2H / tp`` input columns per rank); the draft token table + needs every column. + w3_w1: ``[E, 2 * I / tp, H]``: per expert, this rank's ``I / tp`` rows of up_proj, then + the same rows of gate_proj (the low-M BF16 MoE's layout). + w2: ``[E, H, I / tp]``: per expert, this rank's ``I / tp`` columns of down_proj. + + ``w3_w1`` and ``w2`` hold the slices the built-in MoE loader gives this rank, in plain + row-major layout: the built-in BF16 MoE may keep its own copy in a backend layout. + """ + + eh_proj: torch.Tensor + w3_w1: torch.Tensor + w2: torch.Tensor + + +def load_weights(model, weights, draft: bool) -> Optional[DraftWeights]: + """Load the checkpoint into the target's modules with the built-in loader. + + With ``draft``, first take the MTP layer's ``DraftWeights`` (the built-in loader releases + checkpoint tensors once it has copied them) and return them; None without ``draft`` or + when the checkpoint has no MTP layer. + """ + draft_weights = _draft_weights(model, weights) if draft else None DeepseekV3WeightLoader(model).load_weights(weights) + return draft_weights + + +def _draft_weights(model, weights) -> Optional[DraftWeights]: + """The MTP layer's ``DraftWeights``: checkpoint layer ``num_hidden_layers``.""" + config = model.config + prefix = f"model.layers.{config.num_hidden_layers}" + eh_proj_key = f"{prefix}.eh_proj.weight" + if eh_proj_key not in weights: + return None + mapping = model.model_config.mapping + tp_size, tp_rank = mapping.tp_size, mapping.tp_rank + hidden = config.hidden_size + inter = config.moe_intermediate_size // tp_size + device = model.lm_head.weight.device + eh_proj = weights[eh_proj_key][:].to(device) + w3_w1 = torch.empty( + (config.n_routed_experts, 2 * inter, hidden), dtype=torch.bfloat16, device=device + ) + w2 = torch.empty((config.n_routed_experts, hidden, inter), dtype=torch.bfloat16, device=device) + + def shard(key: str, mode: TensorParallelMode) -> torch.Tensor: + tensor = load_weight_shard(weights[key], tp_size, tp_rank, mode, device=device) + assert tensor.dtype == torch.bfloat16, f"{key}: the MTP layer's experts are not BF16" + return tensor + + for expert in range(config.n_routed_experts): + name = f"{prefix}.mlp.experts.{expert}" + # The built-in MoE loader's slices: up_proj and gate_proj by rows, down_proj by columns. + w3_w1[expert, :inter].copy_(shard(f"{name}.up_proj.weight", TensorParallelMode.COLUMN)) + w3_w1[expert, inter:].copy_(shard(f"{name}.gate_proj.weight", TensorParallelMode.COLUMN)) + w2[expert].copy_(shard(f"{name}.down_proj.weight", TensorParallelMode.ROW)) + return DraftWeights(eh_proj, w3_w1, w2) From 2daccc666fe9f82cbe630da3d617b50524308d69 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 17:52:31 -0700 Subject: [PATCH 41/53] Build the GLM-5.2 tp16 targets without autograd The decode target derives its operands from the loaded parameters, which require grad; under grad mode the derivation kept its intermediates alive. Signed-off-by: Dongfeng Yu --- .../models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py index e8b9a5e0b080..2bd0708e6418 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py @@ -1455,6 +1455,7 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: self._prefill: Optional[PrefillTarget] = None self._decode: Optional[DecodeTarget] = None + @torch.no_grad() def post_load_weights(self) -> None: """Build the targets once the weights are final. Reached more than once (the shell and the loader's module walk), so it only builds once.""" From 2e6cc698930ad2dd71aa2940746cd09869c6d2f1 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 17:55:31 -0700 Subject: [PATCH 42/53] Start the GLM-5.2 tp16 dsa and draft-step test ranks on the imported tensorrt_llm The ranks put the repo root first on their path, so under an installed wheel they imported the source tree, which has no bindings. Signed-off-by: Dongfeng Yu --- .../dsa/_glm_rank_job.py | 15 ++++++--------- .../spec/_draft_step_rank_job.py | 11 +++++------ 2 files changed, 11 insertions(+), 15 deletions(-) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_glm_rank_job.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_glm_rank_job.py index 1f750a9d65b7..ede11eb3e57b 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_glm_rank_job.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_glm_rank_job.py @@ -74,15 +74,12 @@ def run(entry: str) -> None: launcher = Path(__file__).resolve().with_name(f"_{entry}_op_matrix.py") env = dict(os.environ, CUDA_VISIBLE_DEVICES=_devices()) - # The ranks import tensorrt_llm absolutely, and a source checkout is not - # necessarily installed. - # tests/unittest/_torch/modeling_v2/models/glm_moe_dsa//dsa -> repo root. - repo_root = Path(__file__).resolve().parents[8] - assert (repo_root / "tensorrt_llm").is_dir(), ( - f"expected the repo root at {repo_root}, found no tensorrt_llm/ there; " - f"this file moved without its parents[] index following" - ) - env["PYTHONPATH"] = os.pathsep.join(p for p in (str(repo_root), env.get("PYTHONPATH", "")) if p) + # The ranks import tensorrt_llm absolutely: the package this process + # imports, installed or not. + package = importlib.util.find_spec("tensorrt_llm") + assert package is not None and package.origin is not None, "tensorrt_llm is not importable" + package_parent = str(Path(package.origin).resolve().parents[1]) + env["PYTHONPATH"] = os.pathsep.join(p for p in (package_parent, env.get("PYTHONPATH", "")) if p) # Read the budget off the entry rather than restating it: an entry that # raises its own deadline would otherwise be killed by this one first, diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_rank_job.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_rank_job.py index 96277c57c6d5..227701102e53 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_rank_job.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_rank_job.py @@ -41,13 +41,12 @@ def run(entry: str, expected_cells: int) -> None: """Run the matrix for `entry` over WORLD_SIZE devices, expecting it to drive `expected_cells` cells.""" launcher = Path(__file__).resolve().with_name("_draft_step_op_matrix.py") env = dict(os.environ, CUDA_VISIBLE_DEVICES=_devices()) - # The ranks import tensorrt_llm absolutely, and a source checkout is not necessarily installed. - repo_root = Path(__file__).resolve().parents[8] - assert (repo_root / "tensorrt_llm").is_dir(), ( - f"expected the repo root at {repo_root}, found no tensorrt_llm/ there" - ) + # The ranks import tensorrt_llm absolutely: the package this process imports, installed or not. + package = importlib.util.find_spec("tensorrt_llm") + assert package is not None and package.origin is not None, "tensorrt_llm is not importable" + package_parent = str(Path(package.origin).resolve().parents[1]) env["PYTHONPATH"] = os.pathsep.join( - p for p in (str(repo_root), env.get("PYTHONPATH", "")) if p + p for p in (package_parent, env.get("PYTHONPATH", "")) if p ) spec = importlib.util.spec_from_file_location(launcher.stem, launcher) module = importlib.util.module_from_spec(spec) From 84581f6926746a93c9321e4a347aff293b754ceb Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 18:09:43 -0700 Subject: [PATCH 43/53] Allow one BF16 rounding step in the GLM-5.2 draft RMSNorm output bands The band's unit is 2**-8 of the row maximum, and one BF16 ulp of an element in the maximum's binade is up to twice that: a midpoint flip there failed the draft-step matrix's glue cell at 1.02 units. Signed-off-by: Dongfeng Yu --- .../spec/_draft_step_op_matrix.py | 7 ++++--- .../spec/test_modeling_v2_draft_glue_rmsnorm_cat.py | 7 ++++--- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_op_matrix.py index 2dcdc782c683..cf79b26d6da0 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_op_matrix.py @@ -192,14 +192,15 @@ def _compare( The RMSNorm outputs (`normed_out`, `glue_out`) differ from the FP32 reference in the order of the 6144 squares and in the kernel's approximate reciprocal square root: a few 1e-7 relative on the scale, one BF16 - ulp on the rare element that close to a rounding midpoint. So 1 ulp against the row scale and 0.25 ulp of - relative RMS, as for the step-0 glue kernel. + ulp on the rare element that close to a rounding midpoint, which is up to twice the band's unit (2**-8 of the + row maximum) in the maximum's binade. So 2 ulp against the row scale and 0.25 ulp of relative RMS, as for the + step-0 glue kernel. """ for name, ref in expected.items(): out = actual[name] if name in ("normed_out", "glue_out"): try: - assert_within_ulp(out, ref, element_ulp=1.0, rms_ulp=0.25) + assert_within_ulp(out, ref, element_ulp=2.0, rms_ulp=0.25) except AssertionError as exc: raise AssertionError(f"{where}: {name}: {exc}") from exc else: diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py index ee62c3167fe6..4da00457de51 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py @@ -31,10 +31,11 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: reference's correctly rounded one): together a few 1e-7 relative on the scale factor, which moves the BF16 rounding of an output only when it lies that close to a rounding midpoint, by one ulp of that element. - So: 1 ulp element-wise against the row scale, 0.25 ulp of relative RMS. A wrong row, column or weight is off by - the size of the values themselves. + One BF16 ulp of an element in the row maximum's binade is up to twice the band's unit (2**-8 of the row + maximum). So: 2 ulp element-wise against the row scale, 0.25 ulp of relative RMS. A wrong row, column or weight + is off by the size of the values themselves. """ - assert_within_ulp(out, ref, element_ulp=1.0, rms_ulp=0.25) + assert_within_ulp(out, ref, element_ulp=2.0, rms_ulp=0.25) def _build(spec, seed): From 79574e993c3445f4b2655249cd73d4e419ed5815 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 18:14:58 -0700 Subject: [PATCH 44/53] Apply main's pre-commit formatting to the GLM-5.2 tp16 target, its ops, kernels and tests ruff-format, ruff import folding, yapf on the legacy custom-op files, clang-format and cmake-format; the hist top-k kernel's needIn becomes needInGroup for codespell. Signed-off-by: Dongfeng Yu --- .../kernels/glmMoeDsa/glmDenseMlpSmallM.cu | 26 +++---- .../kernels/glmMoeDsa/glmDenseMlpSmallM.h | 4 +- .../kernels/glmMoeDsa/glmDraftSampleGlue.cu | 12 +-- .../kernels/glmMoeDsa/glmDsaHistTopK.cu | 48 ++++++------ .../kernels/glmMoeDsa/glmFusedAGemm.cu | 16 ++-- .../kernels/glmMoeDsa/glmLamport.cuh | 3 +- .../kernels/glmMoeDsa/glmLowMBf16Moe.cu | 31 ++++---- .../kernels/glmMoeDsa/glmMlaKvHalfAppend.cu | 14 ++-- .../kernels/glmMoeDsa/glmMnnvlHop.cu | 22 +++--- .../kernels/glmMoeDsa/glmMnnvlHop.h | 28 +++---- .../kernels/glmMoeDsa/glmRouterGemm.cu | 4 +- .../glmMoeDsa/glmSparseMlaFlashDecode.cu | 25 +++--- .../kernels/moe/trtllmGen/Gemm2K128.cu | 11 ++- .../trtllmGen/routing/RoutingFusedRouter.cu | 32 ++++---- .../kernels/moe/trtllmGen/runner.cu | 4 +- cpp/tensorrt_llm/thop/CMakeLists.txt | 11 ++- .../thop/glmDraftSampleGlueOp.cpp | 12 +-- cpp/tensorrt_llm/thop/glmLowMBf16MoeOp.cpp | 11 ++- .../thop/glmMlaKvHalfAppendOp.cpp | 4 +- cpp/tensorrt_llm/thop/glmMnnvlHopOp.cpp | 4 +- .../thop/glmSparseMlaFlashDecodeOp.cpp | 11 +-- .../glm_5_2_nvfp4__sm_107__tp16/modeling.py | 41 ++++------ .../ops/attention/mla_kv_half_append.py | 12 ++- .../ops/dsa/_fp8_paged_mqa_logits.py | 4 +- .../ops/gemm/_low_m_splitk_fused.py | 24 ++---- .../ops/gemm/indexer_q_tail.py | 7 +- .../ops/gemm/low_m_gemm.py | 7 +- .../ops/moe/routed_experts.py | 12 +-- .../ops/spec/draft_ehproj_table.py | 1 + .../glm_5_2_nvfp4__sm_107__tp16/weights.py | 4 +- .../_torch/custom_ops/cpp_custom_ops.py | 14 ++-- .../custom_ops/trtllm_gen_custom_ops.py | 77 +++++++++---------- .../test_modeling_v2_mla_kv_half_append.py | 4 +- ...est_modeling_v2_sparse_mla_flash_decode.py | 16 ++-- .../comm/_mnnvl_moe_finalize_hop_op_matrix.py | 3 +- .../comm/_rank_job.py | 4 +- .../dsa/_indexer_inputs.py | 4 +- .../gemm/test_modeling_v2_indexer_q_tail.py | 4 +- .../moe/test_modeling_v2_lowm_bf16_moe.py | 10 ++- .../moe/test_modeling_v2_routed_experts.py | 8 +- .../spec/_draft_step_op_matrix.py | 12 +-- .../spec/_draft_step_rank_job.py | 4 +- ...test_modeling_v2_draft_glue_rmsnorm_cat.py | 6 +- .../test_modeling_v2_draft_token_table.py | 6 +- 44 files changed, 279 insertions(+), 338 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.cu index b2178214c4ae..2a5bb29cdad2 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.cu +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.cu @@ -206,8 +206,8 @@ __device__ __forceinline__ void mmaBf16( //! Two MMAs of one step and row tile: MMA A takes contraction elements 0..3 of the thread's 8-element vectors, MMA B //! elements 4..7. Fragments: a0/a2 = weight row "group", a1/a3 = weight row "group + 8", b0/b1 = activation row //! "group". -__device__ __forceinline__ void mmaStep(float (&accA)[4], float (&accB)[4], uint4 const& wLo, uint4 const& wHi, - uint4 const& xv) +__device__ __forceinline__ void mmaStep( + float (&accA)[4], float (&accB)[4], uint4 const& wLo, uint4 const& wHi, uint4 const& xv) { mmaBf16(accA, wLo.x, wHi.x, wLo.y, wHi.y, xv.x, xv.y); mmaBf16(accB, wLo.z, wHi.z, wLo.w, wHi.w, xv.z, xv.w); @@ -256,8 +256,8 @@ __device__ __forceinline__ int weightRow(int tile, int t, int r, int inter) //! x [numRows, kDim], weight [*, kDim]; out: gate/up h [numRows, outDim] (outDim = inter), down out [numRows, outDim]. template -__global__ void __launch_bounds__(T::kThreads, 1) denseMlpSmallMKernel(Bf16 const* __restrict__ x, - Bf16 const* __restrict__ weight, Bf16* __restrict__ out, int numRows, int outDim) +__global__ void __launch_bounds__(T::kThreads, 1) denseMlpSmallMKernel( + Bf16 const* __restrict__ x, Bf16 const* __restrict__ weight, Bf16* __restrict__ out, int numRows, int outDim) { #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) constexpr int kDim = T::kKDim; @@ -432,8 +432,8 @@ __global__ void __launch_bounds__(T::kThreads, 1) denseMlpSmallMKernel(Bf16 cons } else { - asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"( - peerSlot), + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(peerSlot), "r"(__float_as_uint(v0)), "r"(peerBar) : "memory"); } @@ -468,8 +468,8 @@ __global__ void __launch_bounds__(T::kThreads, 1) denseMlpSmallMKernel(Bf16 cons } template -cudaError_t launchTiling(Bf16 const* in, Bf16 const* weight, Bf16* out, int numRows, int outDim, bool enablePdl, - cudaStream_t stream) +cudaError_t launchTiling( + Bf16 const* in, Bf16 const* weight, Bf16* out, int numRows, int outDim, bool enablePdl, cudaStream_t stream) { auto const kernel = &denseMlpSmallMKernel; if constexpr (T::kSmemBytes > 48 * 1024) @@ -521,8 +521,7 @@ bool tilingSupported(int tiling, int numRows, int kDim, int outDim) if (tiling == I) { using T = std::tuple_element_t; - return numRows >= 1 && numRows <= kMaxRows && kDim == T::kKDim && outDim > 0 - && outDim % T::kOutCols == 0; + return numRows >= 1 && numRows <= kMaxRows && kDim == T::kKDim && outDim > 0 && outDim % T::kOutCols == 0; } return tilingSupported(tiling, numRows, kDim, outDim); } @@ -540,8 +539,7 @@ cudaError_t launchTilingIndex(int tiling, Bf16 const* in, Bf16 const* weight, Bf { if (tiling == I) { - return launchTiling>( - in, weight, out, numRows, outDim, enablePdl, stream); + return launchTiling>(in, weight, out, numRows, outDim, enablePdl, stream); } return launchTilingIndex(tiling, in, weight, out, numRows, outDim, enablePdl, stream); } @@ -593,8 +591,8 @@ cudaError_t launchGateUpSwiglu(Bf16 const* x, Bf16 const* gateUpWeight, Bf16* h, return launchTilingIndex(resolved, x, gateUpWeight, h, numRows, inter, enablePdl, stream); } -cudaError_t launchDown(Bf16 const* h, Bf16 const* downWeight, Bf16* out, int numRows, int inter, int hidden, - int tiling, bool enablePdl, cudaStream_t stream) +cudaError_t launchDown(Bf16 const* h, Bf16 const* downWeight, Bf16* out, int numRows, int inter, int hidden, int tiling, + bool enablePdl, cudaStream_t stream) { int const resolved = resolveTiling(tiling, numRows, inter, hidden); if (resolved < 0) diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.h index 90cd1b7f90a6..968c74dc1b55 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.h +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.h @@ -47,8 +47,8 @@ bool isDownSupported(int numRows, int inter, int hidden, int tiling = kDefaultTi //! \param x activations [numRows, hidden], row-major, BF16, 16-byte aligned. //! \param gateUpWeight [2 * inter, hidden], row-major, BF16 (rows [0, inter) gate, [inter, 2 * inter) up). //! \param h output [numRows, inter], row-major, BF16. -cudaError_t launchGateUpSwiglu(__nv_bfloat16 const* x, __nv_bfloat16 const* gateUpWeight, __nv_bfloat16* h, - int numRows, int hidden, int inter, int tiling, bool enablePdl, cudaStream_t stream); +cudaError_t launchGateUpSwiglu(__nv_bfloat16 const* x, __nv_bfloat16 const* gateUpWeight, __nv_bfloat16* h, int numRows, + int hidden, int inter, int tiling, bool enablePdl, cudaStream_t stream); //! Down projection: out[m, n] = bf16(sum_k h[m, k] * W[n, k]) with FP32 sums. //! \param h activations [numRows, inter], row-major, BF16, 16-byte aligned. diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.cu index ea0458c3c894..5a936de3083c 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.cu +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.cu @@ -496,7 +496,8 @@ __global__ void __launch_bounds__(kBlockSize) draftSampleGlueKernel(DraftSampleG { int64_t const token = sTokens[row]; bool const valid = token >= 0 && token < p.vocabSize; - normRowWindow(valid ? reinterpret_cast(p.embedTable + token * kDraftSampleGlueHidden) : nullptr, + normRowWindow( + valid ? reinterpret_cast(p.embedTable + token * kDraftSampleGlueHidden) : nullptr, sWeights[0], p.eps, lo, hi, p.glueOut + row * p.glueOutStride, group, sGroupSums); } } @@ -942,10 +943,11 @@ void invokeDraftSampleGlue(DraftSampleGlueParams const& params, cudaStream_t str || (params.colLo >= 0 && params.colLo < params.colHi && params.colHi <= 2 * kDraftSampleGlueHidden && params.colLo % 8 == 0 && params.colHi % 8 == 0), "draft sample glue: invalid eh_proj input window [%d, %d)", params.colLo, params.colHi); - TLLM_CHECK_WITH_INFO(params.glueOut == nullptr || params.colLo >= kDraftSampleGlueHidden - || params.embedTable != nullptr, + TLLM_CHECK_WITH_INFO( + params.glueOut == nullptr || params.colLo >= kDraftSampleGlueHidden || params.embedTable != nullptr, "draft sample glue: the window needs the embedding table"); - TLLM_CHECK_WITH_INFO(params.glueOut == nullptr || params.colHi <= kDraftSampleGlueHidden || params.hidden != nullptr, + TLLM_CHECK_WITH_INFO( + params.glueOut == nullptr || params.colHi <= kDraftSampleGlueHidden || params.hidden != nullptr, "draft sample glue: the window needs the hidden rows"); TLLM_CHECK_WITH_INFO(params.reqIdxPerToken == nullptr || (params.seqLens != nullptr && params.kvLens != nullptr), "draft sample glue: the DSA metadata needs the sequence and KV lengths"); @@ -989,7 +991,7 @@ void invokeDraftEhprojTable(DraftEhprojTableParams const& params, cudaStream_t s config.numAttrs = 2; // Up to TP8 the kernel keeps eight rank slots (registers of the polled rows); TP16 takes sixteen. auto* const kernel = s.tpSize <= kTableTp8RankSlots ? &draftEhprojTableKernel - : &draftEhprojTableKernel; + : &draftEhprojTableKernel; TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, params)); } diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.cu index aa3f2bee28a9..a73b31395b79 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.cu +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.cu @@ -117,6 +117,7 @@ enum DebugSlot : int32_t kDbgBndCount = 42, kDbgLocalCut = 43 }; + static_assert(kDbgLocalCut < kDebugSlots, "debug slots"); // Shared memory (32-bit words). Peers never read this CTA's shared memory: the cluster's exchanges are pushes @@ -132,14 +133,14 @@ constexpr int32_t kSmemCoarse = kSmemSlots + kBins; // [kMaxRank constexpr int32_t kSmemFine = kSmemCoarse + kMaxRanks * kMaxGroups; // [kMaxRanks][kGroupBins] constexpr int32_t kSmemGsum = kSmemFine + kMaxRanks * kGroupBins; // [2][kMaxGroups] group totals (all, lower ranks) constexpr int32_t kPreRow = kGroupBins + 4; // the group's kGroupBins + 1 suffix sums, padded to 16-byte pushes -constexpr int32_t kSmemPre = kSmemGsum + 2 * kMaxGroups; // [kLevels][kCtas][kPreRow] CTAs' suffix sums +constexpr int32_t kSmemPre = kSmemGsum + 2 * kMaxGroups; // [kLevels][kCtas][kPreRow] CTAs' suffix sums constexpr int32_t kSmemScratch = kSmemPre + kLevels * kCtas * kPreRow; // [256] scan scratch -constexpr int32_t kSmemCut = kSmemScratch + 256; // [32] the cut of the current level -constexpr int32_t kSmemOffs = kSmemCut + 32; // [8] this CTA's offsets (warp 0) -constexpr int32_t kSmemMisc = kSmemOffs + 8; // [16] -constexpr int32_t kSmemSmall = kSmemMisc + 16; // [2 * kSmallRank] -constexpr int32_t kSmemBar = kSmemSmall + 2 * kSmallRank; // [2 * kLevels] uint64 barriers -constexpr int32_t kSmemDbg = kSmemBar + 4 * kLevels; // [kDebugSlots] int64 (debug only) +constexpr int32_t kSmemCut = kSmemScratch + 256; // [32] the cut of the current level +constexpr int32_t kSmemOffs = kSmemCut + 32; // [8] this CTA's offsets (warp 0) +constexpr int32_t kSmemMisc = kSmemOffs + 8; // [16] +constexpr int32_t kSmemSmall = kSmemMisc + 16; // [2 * kSmallRank] +constexpr int32_t kSmemBar = kSmemSmall + 2 * kSmallRank; // [2 * kLevels] uint64 barriers +constexpr int32_t kSmemDbg = kSmemBar + 4 * kLevels; // [kDebugSlots] int64 (debug only) constexpr int32_t kSmemWords = kSmemDbg + 2 * kDebugSlots; constexpr size_t kSmemBytes = size_t{kSmemWords} * sizeof(uint32_t); // The stage area holds, for a cut on the first digit, the level-0 buckets: 16-bit key indices of this CTA (the key at @@ -148,8 +149,7 @@ constexpr int32_t kSmemBucket = kSmemStageArea; // [kKeys * kThreads] uint16 constexpr int32_t kSmemStage = kSmemStageArea; // [kTopK] selected tokens static_assert(kTopK <= kStageWords && kKeys * kThreads <= 65536, "16-bit bucket indices, the stage in the buckets"); static_assert(kSmemKeys % 4 == 0 && kSmemStage % 4 == 0 && kSmemSlots % 4 == 0 && kSmemFine % 4 == 0 - && kSmemPre % 4 == 0 && kSmemScratch % 4 == 0 - && kSmemBar % 2 == 0 && kSmemDbg % 2 == 0, + && kSmemPre % 4 == 0 && kSmemScratch % 4 == 0 && kSmemBar % 2 == 0 && kSmemDbg % 2 == 0, "aligned arrays"); // Below the 132 KB carveout: L1 keeps the register spills of 1024 threads. static_assert(kSmemBytes <= 131 * 1024, "shared memory"); @@ -182,7 +182,7 @@ enum OffsSlot : int32_t enum MiscSlot : int32_t { - kMiscLevelBin = 0, // [kLevels] the cut bin of each level + kMiscLevelBin = 0, // [kLevels] the cut bin of each level kMiscSmallCount = 4, kMiscBucketBin = 5, // the lowest bin of the level-0 buckets kMiscRec = 6, // [3] the previous exchange launch's record of the row (levels, output words, boundary pairs) @@ -347,8 +347,8 @@ __device__ __forceinline__ void fenceBarInit() // 16 bytes into a peer's shared memory, completing on the peer's barrier. __device__ __forceinline__ void pushV4(uint32_t remote, uint4 v, uint32_t remoteBar) { - asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.u32 [%0], {%1, %2, %3, %4}, [%5];" ::"r"( - remote), + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.u32 [%0], {%1, %2, %3, %4}, [%5];" ::"r"(remote), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w), "r"(remoteBar) : "memory"); } @@ -371,8 +371,7 @@ __device__ __forceinline__ int64_t outOff(int32_t buf, int32_t row) __device__ __forceinline__ int64_t bndOff(int32_t buf, int32_t row) { - return int64_t{buf} * kBufferWords + kLevels * kHistRegionWords + kOutRegionWords - + int64_t{row} * kBoundaryCap * 2; + return int64_t{buf} * kBufferWords + kLevels * kHistRegionWords + kOutRegionWords + int64_t{row} * kBoundaryCap * 2; } struct WaitCtx @@ -548,8 +547,8 @@ __device__ __forceinline__ uint32_t blockExclusiveSuffix(uint32_t v, uint32_t* w // flagged in mask a and in mask b (bit 4 k + j): the position of this thread's first flagged key of float4 k is // ba[k] (bb[k]), the next ones follow. Also the block counts of both. One block scan (16-bit fields: a CTA flags at // most 4 keys per thread and float4). -__device__ __forceinline__ void positionBases(uint32_t ma, uint32_t mb, uint32_t* ws, uint32_t (&ba)[kVecs], - uint32_t (&bb)[kVecs], uint32_t& na, uint32_t& nb) +__device__ __forceinline__ void positionBases( + uint32_t ma, uint32_t mb, uint32_t* ws, uint32_t (&ba)[kVecs], uint32_t (&bb)[kVecs], uint32_t& na, uint32_t& nb) { static_assert(kVecs == 3, "three float4 counts per mask"); uint32_t const v[3] = {nibbleCount(ma, 0) | (nibbleCount(ma, 1) << 16), @@ -855,8 +854,8 @@ __device__ void receiveCut(HistTopKParams const& p, uint32_t const* ucLevel, uin int32_t const t = threadIdx.x; int32_t const world = p.worldSize; int32_t const rank = p.rank; - int32_t const ngroups = nbins / kGroupBins; // 64 or 32 - int32_t const vecsPerRank = ngroups / 4; // 16 or 8 + int32_t const ngroups = nbins / kGroupBins; // 64 or 32 + int32_t const vecsPerRank = ngroups / 4; // 16 or 8 uint32_t* const gAllSum = gsum; // [kMaxGroups] all ranks uint32_t* const gLowSum = gsum + kMaxGroups; // [kMaxGroups] the lower ranks if (t == 0) @@ -942,7 +941,7 @@ __device__ void receiveCut(HistTopKParams const& p, uint32_t const* ucLevel, uin // No group holds the need-th key: the ranks' inputs disagree, nothing exact can follow. abortKernel(wc, 200 + what); } - uint32_t const needIn = need - cut[kCutGroupAbove]; + uint32_t const needInGroup = need - cut[kCutGroupAbove]; stamp(dbg, kDbgGroup); // Warps 2..8 (idle in steps 3 and 4): this CTA's suffix sums over the group's bins and the bin above them // (vectors 0..8 of four, the last three words padding) to every CTA of the row, one destination per warp: the 9 @@ -1046,7 +1045,7 @@ __device__ void receiveCut(HistTopKParams const& p, uint32_t const* ucLevel, uin } uint32_t const aAbove = xa - all; uint32_t const lAbove = xl - low; - uint32_t const holder = __ballot_sync(kFull, aAbove < needIn && needIn <= aAbove + all); + uint32_t const holder = __ballot_sync(kFull, aAbove < needInGroup && needInGroup <= aAbove + all); if (holder != 0u) { int32_t const w = lowBit(holder); @@ -1060,7 +1059,7 @@ __device__ void receiveCut(HistTopKParams const& p, uint32_t const* ucLevel, uin } uint32_t const bndRankOff = __reduce_add_sync(kFull, t < rank ? bc : 0u); // When the bin is the exact key: rank r takes what the higher ranks leave of the bin's need. - uint32_t const needBin = needIn - aAboveW; + uint32_t const needBin = needInGroup - aAboveW; uint32_t xs = bc; #pragma unroll for (int32_t o = 1; o < kMaxRanks; o <<= 1) @@ -1091,7 +1090,8 @@ __device__ void receiveCut(HistTopKParams const& p, uint32_t const* ucLevel, uin stamp(dbg, kDbgCut); // When this is the final cut (as the main loop decides it), this CTA's offsets at once. uint32_t const needAfter = need - aboveAll; - int32_t const mode = needAfter == allW ? kModeFull + int32_t const mode = needAfter == allW + ? kModeFull : (level == kLevels - 1 ? kModeTies : (allW <= static_cast(kBoundaryCap) ? kModeCand : kModeNext)); if (mode != kModeNext) @@ -1954,8 +1954,8 @@ void invokeDsaHistShardMeta(int32_t const* kvLens, int32_t const* blockTable, in { throw std::invalid_argument("glm_dsa_hist_shard_meta: unsupported group or page size"); } - dsaHistShardMetaKernel<<<1, 256, 0, stream>>>(kvLens, blockTable, blockTableWidth, worldSize, rank, - tokensPerBlock, shardCtx, shardBlockTable, shardBlockTableWidth); + dsaHistShardMetaKernel<<<1, 256, 0, stream>>>(kvLens, blockTable, blockTableWidth, worldSize, rank, tokensPerBlock, + shardCtx, shardBlockTable, shardBlockTableWidth); cudaError_t const err = cudaGetLastError(); if (err != cudaSuccess) { diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.cu index 3f8d4b4a959f..5224fcb38651 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.cu +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.cu @@ -107,8 +107,8 @@ __device__ void prefetch_ranges_l2(FusedAL2Prefetch const& ranges, int local_tid // Tensor copies completing on an mbarrier, ldmatrix and HMMA on shared addresses. __device__ __forceinline__ void mma8_expect_tx(uint64_t* smem_barrier, uint32_t bytes) { - asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"( - __nvvm_get_smem_pointer(smem_barrier)), + asm volatile( + "mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(__nvvm_get_smem_pointer(smem_barrier)), "r"(bytes) : "memory"); } @@ -252,8 +252,8 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_mma8_kernel(bf16_t* outpu #pragma unroll for (int j = 0; j < 2; j++) { - mma8_tma_load_3d(smem_b + (2 * warp + j) * a_tile_bytes, &tma_a, 0, m0, - (warp + warp_cnt * (j + 1)) * boxes, staged + warp); + mma8_tma_load_3d(smem_b + (2 * warp + j) * a_tile_bytes, &tma_a, 0, m0, (warp + warp_cnt * (j + 1)) * boxes, + staged + warp); } } uint4 a_reg[2][pair_cnt][2]; @@ -424,8 +424,8 @@ void launchFusedAMma8(bf16_t* output, bf16_t const* hiddenStates, bf16_t const* CUtensorMap const tmaA = makeFusedAMma8Map(weight, gemm_m, kFusedAHiddenSize, kFusedARowsPerCta, true); CUtensorMap const tmaB = makeFusedAMma8Map(hiddenStates, numTokens, kFusedAHiddenSize, numTokens, numTokens == kFusedAMaxTokens); - auto* kernel = l2Prefetch != nullptr ? &fused_a_gemm_mma8_kernel - : &fused_a_gemm_mma8_kernel; + auto* kernel + = l2Prefetch != nullptr ? &fused_a_gemm_mma8_kernel : &fused_a_gemm_mma8_kernel; TLLM_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes)); cudaLaunchConfig_t config; config.gridDim = dim3(gemm_m / kFusedARowsPerCta, 1, 1); @@ -437,8 +437,8 @@ void launchFusedAMma8(bf16_t* output, bf16_t const* hiddenStates, bf16_t const* attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); config.numAttrs = 1; config.attrs = attrs; - TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, output, numTokens, - l2Prefetch != nullptr ? *l2Prefetch : FusedAL2Prefetch{}, tmaA, tmaB)); + TLLM_CUDA_CHECK(cudaLaunchKernelEx( + &config, kernel, output, numTokens, l2Prefetch != nullptr ? *l2Prefetch : FusedAL2Prefetch{}, tmaA, tmaB)); } } // namespace diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLamport.cuh b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLamport.cuh index c7461c9dbc00..a5147530112e 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLamport.cuh +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLamport.cuh @@ -98,8 +98,7 @@ public: for (uint32_t stage = 0; stage < mDirtyBufferLayout.numStages; stage++) { uint32_t const fillBoundary = (mBytesToClear[stage] + sizeof(PackedType) - 1) / sizeof(PackedType); - auto* dirty - = static_cast(mDirtyBufferLayout.getStagePtr(bufferBasePtr, mDirtyIndex, stage)); + auto* dirty = static_cast(mDirtyBufferLayout.getStagePtr(bufferBasePtr, mDirtyIndex, stage)); for (uint32_t packedIdx = globalTid; packedIdx < fillBoundary; packedIdx += numThreads) { dirty[packedIdx] = fillValue; diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.cu index 850a0c2bcd7f..bc67707178a8 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.cu +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.cu @@ -372,9 +372,9 @@ __device__ __forceinline__ void warpSum2(float& a, float& b) // the two overlap; each consumer warp reduces its shares with shuffles into shared memory, and each unit's SwiGLU adds // the 8 warp sums in a fixed order. template -__global__ void __launch_bounds__(kFc1Threads, 2) lowMBf16MoeFc1Kernel(Bf16* __restrict__ act, - Bf16 const* __restrict__ x, float const* __restrict__ routerLogits, BiasT const* __restrict__ bias, - Bf16 const* __restrict__ w3w1, int hiddenSize, int interSize, int numStages) +__global__ void __launch_bounds__(kFc1Threads, 2) + lowMBf16MoeFc1Kernel(Bf16* __restrict__ act, Bf16 const* __restrict__ x, float const* __restrict__ routerLogits, + BiasT const* __restrict__ bias, Bf16 const* __restrict__ w3w1, int hiddenSize, int interSize, int numStages) { #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) extern __shared__ __align__(128) uint8_t smem[]; @@ -490,8 +490,8 @@ __global__ void __launch_bounds__(kFc1Threads, 2) lowMBf16MoeFc1Kernel(Bf16* __r { if (i < piecesPerThread) { - unpackPiece(__ldcg(reinterpret_cast(x) + threadIdx.x + i * kFc1ConsumerThreads), - xValues[i]); + unpackPiece( + __ldcg(reinterpret_cast(x) + threadIdx.x + i * kFc1ConsumerThreads), xValues[i]); } } // Wait at the barrier, without polling the ring, while the producer selects the experts and issues the first @@ -574,9 +574,9 @@ __global__ void __launch_bounds__(kFc1Threads, 2) lowMBf16MoeFc1Kernel(Bf16* __r // that wait, from the router logits and the bias staged by bulk copies, without FP64 instructions. After the wait every // warp stages its expert's SwiGLU outputs itself. template -__global__ void __launch_bounds__(kFc2Threads, 2) lowMBf16MoeFc2Kernel(Bf16* __restrict__ out, - Bf16 const* __restrict__ act, float const* __restrict__ routerLogits, BiasT const* __restrict__ bias, - Bf16 const* __restrict__ w2, int hiddenSize, int interSize, RoutingScale scale) +__global__ void __launch_bounds__(kFc2Threads, 2) + lowMBf16MoeFc2Kernel(Bf16* __restrict__ out, Bf16 const* __restrict__ act, float const* __restrict__ routerLogits, + BiasT const* __restrict__ bias, Bf16 const* __restrict__ w2, int hiddenSize, int interSize, RoutingScale scale) { #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) constexpr int kRows = kTiles * kMmaRows; @@ -712,8 +712,7 @@ __global__ void __launch_bounds__(kFc2Threads, 2) lowMBf16MoeFc2Kernel(Bf16* __r for (int tileIdx = 0; tileIdx < kTiles; ++tileIdx) { uint4 const lo = *reinterpret_cast(slice + (tileIdx * kMmaRows + g) * interSize + j); - uint4 const hi - = *reinterpret_cast(slice + (tileIdx * kMmaRows + g + 8) * interSize + j); + uint4 const hi = *reinterpret_cast(slice + (tileIdx * kMmaRows + g + 8) * interSize + j); mmaStep(acc[tileIdx][half], lo, hi, av); } } @@ -773,8 +772,8 @@ RoutingScale makeRoutingScale(double routedScalingFactor) } template -void launchFc2(Bf16* out, Bf16 const* act, float const* routerLogits, BiasT const* bias, Bf16 const* w2, - int hiddenSize, int interSize, double routedScalingFactor, cudaStream_t stream) +void launchFc2(Bf16* out, Bf16 const* act, float const* routerLogits, BiasT const* bias, Bf16 const* w2, int hiddenSize, + int interSize, double routedScalingFactor, cudaStream_t stream) { int const rows = kTiles * kMmaRows; size_t const smemBytes = static_cast(kTopK) * (rows + 1) * interSize * sizeof(Bf16); @@ -806,12 +805,12 @@ void invokeLowMBf16Moe(Bf16* out, Bf16* act, Bf16 const* x, float const* routerL int const rowBytes = hiddenSize * static_cast(sizeof(Bf16)); int const numStages = std::min(kMaxStages, kFc1RingBytes / rowBytes); int const maxRows = 2 * ((numUnits + fc1Grid - 1) / fc1Grid); - size_t const fc1SmemBytes = static_cast(numStages) * rowBytes - + static_cast(maxRows) * kFc1ConsumerWarps * sizeof(float); + size_t const fc1SmemBytes + = static_cast(numStages) * rowBytes + static_cast(maxRows) * kFc1ConsumerWarps * sizeof(float); auto* fc1 = &lowMBf16MoeFc1Kernel; setSharedMemoryAttributes(fc1, fc1SmemBytes); - common::launchWithPdlWhenEnabled("lowMBf16MoeFc1", fc1, dim3(fc1Grid), dim3(kFc1Threads), fc1SmemBytes, stream, - act, x, routerLogits, bias, w3w1, hiddenSize, interSize, numStages); + common::launchWithPdlWhenEnabled("lowMBf16MoeFc1", fc1, dim3(fc1Grid), dim3(kFc1Threads), fc1SmemBytes, stream, act, + x, routerLogits, bias, w3w1, hiddenSize, interSize, numStages); // Two 16-row tiles per FC2 CTA where the 8 weight slices fit in shared memory, else one. if (interSize <= kMaxInterSize / 2) diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.cu index 43513de09ad2..91246637cf1f 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.cu +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.cu @@ -67,11 +67,11 @@ inline __device__ uint2 toFp8x8(float const (&v)[kEltsPerVec]) // kTriggerAtEntry: the PDL dependents may launch at kernel entry instead of after the stores; they read this kernel's // outputs only after their own grid-dependency wait. template -__global__ void __launch_bounds__(kBlockThreads) glmMlaKvHalfAppendKernel(__nv_bfloat16 const* __restrict__ latent, - int latentRowStride, KVBlockArray kvCache, float2 const* __restrict__ cosSinCache, - __nv_bfloat16 const* __restrict__ kvNormWeight, float kvNormEps, int numTokens, int seqLenQ, - int32_t const* __restrict__ kvLens, float* __restrict__ bmm1Scale, float* __restrict__ bmm2Scale, - float softmaxScale) +__global__ void __launch_bounds__(kBlockThreads) + glmMlaKvHalfAppendKernel(__nv_bfloat16 const* __restrict__ latent, int latentRowStride, KVBlockArray kvCache, + float2 const* __restrict__ cosSinCache, __nv_bfloat16 const* __restrict__ kvNormWeight, float kvNormEps, + int numTokens, int seqLenQ, int32_t const* __restrict__ kvLens, float* __restrict__ bmm1Scale, + float* __restrict__ bmm2Scale, float softmaxScale) { constexpr int kRowsPerBlock = kBlockThreads / kWarpSize; int const lane = threadIdx.x % kWarpSize; @@ -219,8 +219,8 @@ void invokeGlmMlaKvHalfAppend(__nv_bfloat16 const* latent, int32_t latentRowStri constexpr int kRowsPerBlock = kBlockSize / kWarpSize; dim3 const grid(static_cast(tensorrt_llm::common::divUp(std::max(numTokens, 1), kRowsPerBlock))); - auto const kernel = triggerAtEntry ? glmMlaKvHalfAppendKernel - : glmMlaKvHalfAppendKernel; + auto const kernel + = triggerAtEntry ? glmMlaKvHalfAppendKernel : glmMlaKvHalfAppendKernel; tensorrt_llm::common::launchWithPdlWhenEnabled("glmMlaKvHalfAppend", kernel, grid, kBlockSize, 0, stream, latent, static_cast(latentRowStride), kvCache, cosSinCache, kvNormWeight, kvNormEps, static_cast(numTokens), static_cast(seqLenQ), kvLens, bmm1Scale, bmm2Scale, softmaxScale); diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.cu index 41e9f7009465..637fa3c8e6f4 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.cu +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.cu @@ -266,15 +266,15 @@ __device__ __forceinline__ float clusterSumOfSquaresFast(float threadSum) if (laneId < numBlocks) { // Lane l sends to CTA l. - uint32_t const slot = static_cast(__cvta_generic_to_shared( - reinterpret_cast(storage.warpSums) + blockRank * kMaxWarps + warpId)); + uint32_t const slot = static_cast( + __cvta_generic_to_shared(reinterpret_cast(storage.warpSums) + blockRank * kMaxWarps + warpId)); uint32_t peerSlot; uint32_t peerBarrier; asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(peerSlot) : "r"(slot), "r"(laneId)); asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(peerBarrier) : "r"(barrier), "r"(laneId)); asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(peerSlot), - "r"(__float_as_uint(warpSum)), "r"(peerBarrier) - : "memory"); + "r"(__float_as_uint(warpSum)), "r"(peerBarrier) + : "memory"); } uint32_t complete = 0; while (!complete) @@ -379,9 +379,9 @@ __device__ __forceinline__ void hopEpilogue(HopKernelParams const& params, GlmLa #pragma unroll for (int i = 0; i < kEltsPerThread; i++) { - packedAccum.elements[i] = cuda_cast<__nv_bfloat16, float>( - cuda_cast(packedAccum.elements[i]) * rcpRms - * cuda_cast(gamma.elements[i])); + packedAccum.elements[i] + = cuda_cast<__nv_bfloat16, float>(cuda_cast(packedAccum.elements[i]) * rcpRms + * cuda_cast(gamma.elements[i])); } reinterpret_cast(¶ms.normOut[threadOffset])[0] = packedAccum.packed; if (rotationPending) @@ -508,8 +508,8 @@ __global__ void __launch_bounds__(kMaxThreads) mnnvlHopKernel(HopKernelParams pa // grid launches because FC1 and FC2 trigger their dependents only after their own waits. The launch bound lets the // eight FC2 rows of a thread be in flight in registers. template -__global__ void __launch_bounds__(kMaxThreads) mnnvlMoeFinalizeHopKernel( - HopKernelParams params, MoeFinalizeKernelParams moe) +__global__ void __launch_bounds__(kMaxThreads) + mnnvlMoeFinalizeHopKernel(HopKernelParams params, MoeFinalizeKernelParams moe) { #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) namespace cg = cooperative_groups; @@ -610,8 +610,8 @@ __global__ void __launch_bounds__(kMaxThreads) mnnvlMoeFinalizeHopKernel( // ==================== Reduction ==================== bool rotationPending = threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0; uint32_t arrivals = 0; - PackedBf16x8 const packedAccum = reduceOneshotPollAll(val, params.rank, stagePtrLocal, - token, params.tokenDim, packedIdx, rotationPending ? flag.arrivalCounter() : nullptr, &arrivals); + PackedBf16x8 const packedAccum = reduceOneshotPollAll(val, params.rank, stagePtrLocal, token, + params.tokenDim, packedIdx, rotationPending ? flag.arrivalCounter() : nullptr, &arrivals); uint4 const bytesToClear{ static_cast(params.numTokens * params.tokenDim * WorldSize * sizeof(__nv_bfloat16)), 0, 0, 0}; rotateIfAllArrived(flag, rotationPending, arrivals, bytesToClear); diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.h index 693646e3eade..aaa7d02be3af 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.h +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.h @@ -35,20 +35,20 @@ constexpr int kMnnvlHopMoeTopK = 8; //! value; the workspace of the one-shot MNNVL all-reduce). struct MnnvlHopParams { - int nRanks; //!< Ranks of the group: 2, 4, 8 or 16 - int rank; //!< This rank - int numTokens; //!< Rows - int tokenDim; //!< Elements per row - void** bufferPtrsDev; //!< Unicast addresses of every rank's workspace (a device array) - void* multicastPtr; //!< Multicast address of the workspace - uint32_t* bufferFlags; //!< Lamport flags of the workspace - void const* input; //!< This rank's partial rows [numTokens, tokenDim]; not read when shardPublished - void const* residualIn; //!< [numTokens, tokenDim] - void const* gamma; //!< RMSNorm weight [tokenDim] - float epsilon; //!< RMSNorm epsilon - void* residualOut; //!< Rows summed over the ranks plus residualIn [numTokens, tokenDim] - void* normOut; //!< RMSNorm of residualOut times gamma [numTokens, tokenDim] - bool shardPublished; //!< The producer of input already stored this rank's rows into the workspace + int nRanks; //!< Ranks of the group: 2, 4, 8 or 16 + int rank; //!< This rank + int numTokens; //!< Rows + int tokenDim; //!< Elements per row + void** bufferPtrsDev; //!< Unicast addresses of every rank's workspace (a device array) + void* multicastPtr; //!< Multicast address of the workspace + uint32_t* bufferFlags; //!< Lamport flags of the workspace + void const* input; //!< This rank's partial rows [numTokens, tokenDim]; not read when shardPublished + void const* residualIn; //!< [numTokens, tokenDim] + void const* gamma; //!< RMSNorm weight [tokenDim] + float epsilon; //!< RMSNorm epsilon + void* residualOut; //!< Rows summed over the ranks plus residualIn [numTokens, tokenDim] + void* normOut; //!< RMSNorm of residualOut times gamma [numTokens, tokenDim] + bool shardPublished; //!< The producer of input already stored this rank's rows into the workspace cudaStream_t stream; }; diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.cu index c6196047a3b2..1b5a83166840 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.cu +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.cu @@ -203,8 +203,8 @@ void launchRouterGemm( config.numAttrs = 1; config.attrs = attrs; TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, - router_gemm_kernel, logits, - hiddenStates, weight)); + router_gemm_kernel, logits, hiddenStates, + weight)); } // Launches the instance of numTokens, one of kNumTokens..kRouterMaxTokens. diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.cu index 795a11fcbce2..af9aade1887d 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.cu +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.cu @@ -144,8 +144,8 @@ __device__ __forceinline__ void mbarrierWaitCluster(uint32_t bar, uint32_t parit __device__ __forceinline__ void bulkCopyGlobalToShared(uint32_t dst, void const* src, uint32_t bytes, uint32_t bar) { asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n" ::"r"(dst), - "l"(src), "r"(bytes), "r"(bar) - : "memory"); + "l"(src), "r"(bytes), "r"(bar) + : "memory"); } // Two floats into the shared memory of a CTA of the cluster; the bytes complete on that CTA's barrier. @@ -160,11 +160,10 @@ __device__ __forceinline__ void storeRemote(uint32_t remoteAddr, float a, float // Sixteen bytes into the shared memory of a CTA of the cluster; the bytes complete on that CTA's barrier. __device__ __forceinline__ void storeRemote16(uint32_t remoteAddr, uint4 value, uint32_t remoteBar) { - asm volatile( - "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];\n" ::"r"( - remoteAddr), - "r"(value.x), "r"(value.y), "r"(value.z), "r"(value.w), "r"(remoteBar) - : "memory"); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];\n" ::"r"( + remoteAddr), + "r"(value.x), "r"(value.y), "r"(value.z), "r"(value.w), "r"(remoteBar) + : "memory"); } __device__ __forceinline__ uint32_t mapToRank(uint32_t addr, uint32_t rank) @@ -249,8 +248,8 @@ __device__ __forceinline__ void mmaE4m3AsF16( float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t const (&bh)[4]) { mmaF16(d, e4m3x2ToHalf2(a0), e4m3x2ToHalf2(a1), e4m3x2ToHalf2(a2), e4m3x2ToHalf2(a3), bh[0], bh[1]); - mmaF16(d, e4m3x2ToHalf2(a0 >> 16), e4m3x2ToHalf2(a1 >> 16), e4m3x2ToHalf2(a2 >> 16), e4m3x2ToHalf2(a3 >> 16), - bh[2], bh[3]); + mmaF16(d, e4m3x2ToHalf2(a0 >> 16), e4m3x2ToHalf2(a1 >> 16), e4m3x2ToHalf2(a2 >> 16), e4m3x2ToHalf2(a3 >> 16), bh[2], + bh[3]); } __device__ __forceinline__ uint8_t toE4m3(float x) @@ -280,8 +279,8 @@ struct EarlyGatherRow // Lane l < 16 of the warp maps index slot `slot` of query token `tokenQ` (a position) to its pool row with the remap // arithmetic, and bulk-copies the row unless it holds one of this step's tokens. Masked lanes copy row 0, as after the // wait. Every read here (indices, block table, KV length, rows of earlier steps) precedes this grid's launch. -__device__ __forceinline__ EarlyGatherRow gatherRowBeforeWait(GlmSparseMlaFlashDecodeParams const& params, int tokenQ, - int slot, int lane, uint32_t dst, uint32_t barGather) +__device__ __forceinline__ EarlyGatherRow gatherRowBeforeWait( + GlmSparseMlaFlashDecodeParams const& params, int tokenQ, int slot, int lane, uint32_t dst, uint32_t barGather) { int position = -1; if (lane < kTokensPerWarp && slot < params.topK) @@ -311,8 +310,8 @@ __device__ __forceinline__ EarlyGatherRow gatherRowBeforeWait(GlmSparseMlaFlashD __syncwarp(); if (lane < kTokensPerWarp && !out.fresh) { - bulkCopyGlobalToShared(dst, static_cast(params.kvPool) + out.row * kRowBytes, kRowBytes, - barGather); + bulkCopyGlobalToShared( + dst, static_cast(params.kvPool) + out.row * kRowBytes, kRowBytes, barGather); } return out; } diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/Gemm2K128.cu b/cpp/tensorrt_llm/kernels/moe/trtllmGen/Gemm2K128.cu index 651cb992ec02..57b00b763d83 100644 --- a/cpp/tensorrt_llm/kernels/moe/trtllmGen/Gemm2K128.cu +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/Gemm2K128.cu @@ -136,8 +136,8 @@ __device__ __forceinline__ void tmaLoad2d(void* dst, void const* map, int x, int __device__ __forceinline__ void bulkLoad(void* dst, void const* src, uint32_t bytes, uint64_t* bar) { - asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( - smemU32(dst)), + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"(smemU32(dst)), "l"(src), "r"(bytes), "r"(smemU32(bar)) : "memory"); } @@ -539,8 +539,7 @@ __global__ void __launch_bounds__(kThreads, 1) gemm2K128Kernel(const __grid_cons { if (j < numSlots) { - *reinterpret_cast(smemAct + j * kActBytes + r8 * 64 + ((cnk ^ ((r8 >> 1) & 3)) << 4)) - = actV[j]; + *reinterpret_cast(smemAct + j * kActBytes + r8 * 64 + ((cnk ^ ((r8 >> 1) & 3)) << 4)) = actV[j]; *reinterpret_cast(smemSfb + j * kSfbBytes + lane * 16) = make_uint4(sf0[j], 0u, sf1[j], 0u); } } @@ -620,8 +619,8 @@ void encodeWeightMap(CUtensorMap* map, void const* weights, uint64_t rows, uint3 bool isSupported(int32_t hiddenSize, int32_t maxNumCtas) { int const sm = tensorrt_llm::common::getSMVersion(); - if (sm < 100 || sm >= 110 || hiddenSize % kTileRows != 0 || hiddenSize / kTileRows < kMaxUnits - || maxNumCtas < 1 || maxNumCtas > kMaxTokenTiles) + if (sm < 100 || sm >= 110 || hiddenSize % kTileRows != 0 || hiddenSize / kTileRows < kMaxUnits || maxNumCtas < 1 + || maxNumCtas > kMaxTokenTiles) { return false; } diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingFusedRouter.cu b/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingFusedRouter.cu index 329225b1021f..7e447cdceb4a 100644 --- a/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingFusedRouter.cu +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingFusedRouter.cu @@ -59,15 +59,15 @@ constexpr int kNumExperts = 256; constexpr int kTopK = 8; constexpr int kMaxTokens = 8; constexpr int kHiddenDim = 6144; -constexpr int kEltsPerVec = 8; // bf16 per 16-byte load, as in the router GEMV -constexpr int kKIterElts = kEltsPerVec * 128; // elements per k-iteration of the router GEMV -constexpr int kKIters = kHiddenDim / kKIterElts; // 6 -constexpr int kSlices = 4; // warps per expert in the router GEMV +constexpr int kEltsPerVec = 8; // bf16 per 16-byte load, as in the router GEMV +constexpr int kKIterElts = kEltsPerVec * 128; // elements per k-iteration of the router GEMV +constexpr int kKIters = kHiddenDim / kKIterElts; // 6 +constexpr int kSlices = 4; // warps per expert in the router GEMV constexpr int kExpertsPerCta = 8; -constexpr int kNumGemvCtas = kNumExperts / kExpertsPerCta * kSlices; // 128 +constexpr int kNumGemvCtas = kNumExperts / kExpertsPerCta * kSlices; // 128 constexpr int kNumThreads = 256; -constexpr int kSliceVecs = kKIters * WarpSize; // 16-byte vectors of one token's slice -constexpr int kActBytes = kMaxTokens * kSliceVecs * 16; // 24 KB +constexpr int kSliceVecs = kKIters * WarpSize; // 16-byte vectors of one token's slice +constexpr int kActBytes = kMaxTokens * kSliceVecs * 16; // 24 KB constexpr int kPartBytes = WarpSize * kExpertsPerCta * kMaxTokens * 4; // 8 KB constexpr int kSmemBytes = kActBytes + kPartBytes; constexpr int kExchangeFloats = kMaxTokens * kNumExperts * kSlices; @@ -100,9 +100,7 @@ __device__ __forceinline__ void gridDependencyWait() __device__ __forceinline__ uint4 loadWeight16(void const* ptr) { uint4 v; - asm volatile("ld.global.nc.v4.u32 {%0, %1, %2, %3}, [%4];" - : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) - : "l"(ptr)); + asm volatile("ld.global.nc.v4.u32 {%0, %1, %2, %3}, [%4];" : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(ptr)); return v; } @@ -398,9 +396,7 @@ __global__ void __launch_bounds__(kNumThreads, 1) routingFusedRouterKernel(Fused return rank; }; auto tileRows = [&](int32_t n) - { - return params.mIsPow2 ? mulLog2(n, params.mPaddingLog2) : mulTileN(n, params.mTileTokensDim); - }; + { return params.mIsPow2 ? mulLog2(n, params.mPaddingLog2) : mulTileN(n, params.mTileTokensDim); }; { int32_t const expert = threadIdx.x; @@ -455,8 +451,8 @@ __global__ void __launch_bounds__(kNumThreads, 1) routingFusedRouterKernel(Fused cudaTriggerProgrammaticLaunchCompletion(); } #else -__global__ void __launch_bounds__(kNumThreads, 1) routingFusedRouterKernel( - FusedRouterParams const __grid_constant__ /* p */) +__global__ void __launch_bounds__(kNumThreads, 1) + routingFusedRouterKernel(FusedRouterParams const __grid_constant__ /* p */) { assert(false && "routingFusedRouterKernel is only supported on SM100+ architectures"); } @@ -482,9 +478,9 @@ void launchFusedRouter(FusedRouterParams const& params, bool usePdl, cudaStream_ bool supportsFusedRouter(Data const& data, int32_t hiddenDim) { - return hiddenDim == kHiddenDim && data.mNumExperts == kNumExperts && data.mTopK == kTopK - && data.mNumTokens >= 1 && data.mNumTokens <= kMaxTokens && data.mNumTokens <= data.mTileTokensDim - && data.mLocalExpertsStrideLog2 == 0 && data.mPreprocessType == RoutingPreprocessType::SigmoidBias + return hiddenDim == kHiddenDim && data.mNumExperts == kNumExperts && data.mTopK == kTopK && data.mNumTokens >= 1 + && data.mNumTokens <= kMaxTokens && data.mNumTokens <= data.mTileTokensDim && data.mLocalExpertsStrideLog2 == 0 + && data.mPreprocessType == RoutingPreprocessType::SigmoidBias && data.mPostprocessType == RoutingPostprocessType::ScaledSumNormalize && data.mDtypeOutput == tg::Dtype::Bfloat16 && data.mPtrTopKWeights != nullptr && data.mPtrRoutingBias != nullptr && data.mPtrTopKIds == nullptr; diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.cu b/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.cu index cecc50daa877..fa2c2d9b6327 100644 --- a/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.cu +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.cu @@ -916,8 +916,8 @@ void Runner::run( = Routing::getMaxNumCgasInBatchDim(args.num_tokens, totalExpertsPerToken, totalLocalExperts, mTileTokensDim); // gemm2K128 covers the FC2 of the NVFP4 MoE at intermediate size 128 and tile N 8 (for example GLM-5 decode at MoE // TP16): the same output rows as the Gemm2 cubin, bitwise (see Gemm2K128.cu). - if (args.prefer_gemm2_k128 && mDtypeAct == btg::Dtype::E2m1 && mDtypeWeights == btg::Dtype::E2m1 - && !mUseDeepSeekFp8 && !mUseFineGrained && mTileTokensDim == 8 && args.num_fused_shared_experts == 0 + if (args.prefer_gemm2_k128 && mDtypeAct == btg::Dtype::E2m1 && mDtypeWeights == btg::Dtype::E2m1 && !mUseDeepSeekFp8 + && !mUseFineGrained && mTileTokensDim == 8 && args.num_fused_shared_experts == 0 && args.intermediate_size == 128 && args.valid_intermediate_size.value_or(args.intermediate_size) == args.intermediate_size && args.valid_hidden_size.value_or(args.hidden_size) == gemm2HiddenSize && args.gemm2_bias == nullptr diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index b649ddc4a283..c4dd5613839e 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -121,9 +121,14 @@ add_th_common_family( # TODO: th_common_moe does not compile with USING_OSS_CUTLASS_MOE_GEMM=OFF # because moe/utils/moeUtilOp.cpp depends on OSS-specific MoE utility types and # function signatures. -add_th_common_family(th_common_moe fusedTopkSoftmax.cpp llama4MinLatency.cpp - noAuxTcOp.cpp dsv3RouterGemmOp.cpp glmRouterGemmOp.cpp - glmLowMBf16MoeOp.cpp) +add_th_common_family( + th_common_moe + fusedTopkSoftmax.cpp + llama4MinLatency.cpp + noAuxTcOp.cpp + dsv3RouterGemmOp.cpp + glmRouterGemmOp.cpp + glmLowMBf16MoeOp.cpp) # MoE torch ops live in moe//; that directory adds them to # th_common_moe with target_sources() so the list stays next to the files. diff --git a/cpp/tensorrt_llm/thop/glmDraftSampleGlueOp.cpp b/cpp/tensorrt_llm/thop/glmDraftSampleGlueOp.cpp index 27390af30f41..b9bbcadff112 100644 --- a/cpp/tensorrt_llm/thop/glmDraftSampleGlueOp.cpp +++ b/cpp/tensorrt_llm/thop/glmDraftSampleGlueOp.cpp @@ -60,11 +60,11 @@ int64_t setDraftStepExchange(tk::DraftSampleGlueParams& params, char const* op, TORCH_CHECK(logits.size(1) % 8 == 0 && isAligned16(logits), op, ": logits rows must be 16-byte aligned and a multiple of 8 columns"); int64_t const bufferBytes = exchange_buffer.numel() * exchange_buffer.element_size() / 3; - TORCH_CHECK(bufferBytes >= rows * group_size * slotBytes, op, ": the Lamport buffer is too small for ", rows, - " rows"); + TORCH_CHECK( + bufferBytes >= rows * group_size * slotBytes, op, ": the Lamport buffer is too small for ", rows, " rows"); CHECK_INPUT(draft_tokens, at::ScalarType::Int); - TORCH_CHECK(draft_tokens.dim() == 2 && draft_tokens.size(0) == rows && draft_step >= 0 - && draft_step < draft_tokens.size(1), + TORCH_CHECK( + draft_tokens.dim() == 2 && draft_tokens.size(0) == rows && draft_step >= 0 && draft_step < draft_tokens.size(1), op, ": draft_tokens must be [rows, draft_len] with draft_step inside"); params.logits = reinterpret_cast<__nv_bfloat16 const*>(logits.data_ptr()); @@ -211,8 +211,8 @@ void glmDraftSampleGlue(at::Tensor const& logits, at::Tensor& exchange_buffer, a for (at::Tensor const& weight : {enorm_weight.value(), hnorm_weight.value()}) { CHECK_INPUT(weight, at::ScalarType::BFloat16); - TORCH_CHECK(weight.numel() == kHidden && isAligned16(weight), - "glm_draft_sample_glue: norm weights must be [6144]"); + TORCH_CHECK( + weight.numel() == kHidden && isAligned16(weight), "glm_draft_sample_glue: norm weights must be [6144]"); } TORCH_CHECK(glue_out->scalar_type() == at::ScalarType::BFloat16 && glue_out->dim() == 2 && glue_out->size(0) == rows && glue_out->size(1) == col_hi - col_lo && glue_out->stride(1) == 1 diff --git a/cpp/tensorrt_llm/thop/glmLowMBf16MoeOp.cpp b/cpp/tensorrt_llm/thop/glmLowMBf16MoeOp.cpp index 383d6c0cb676..e414e334d7b8 100644 --- a/cpp/tensorrt_llm/thop/glmLowMBf16MoeOp.cpp +++ b/cpp/tensorrt_llm/thop/glmLowMBf16MoeOp.cpp @@ -53,8 +53,7 @@ at::Tensor glmLowmBf16Moe(at::Tensor const& x, at::Tensor const& routerLogits, a CHECK_INPUT(w2, at::ScalarType::BFloat16); TORCH_CHECK(eScoreCorrectionBias.scalar_type() == at::ScalarType::Float || eScoreCorrectionBias.scalar_type() == at::ScalarType::BFloat16, - "glm_lowm_bf16_moe expects an FP32 or BF16 e_score_correction_bias, got ", - eScoreCorrectionBias.scalar_type()); + "glm_lowm_bf16_moe expects an FP32 or BF16 e_score_correction_bias, got ", eScoreCorrectionBias.scalar_type()); for (at::Tensor const* t : {&routerLogits, &eScoreCorrectionBias, &w3W1, &w2}) { TORCH_CHECK(t->device() == x.device(), "glm_lowm_bf16_moe: all tensors must be on the device of x"); @@ -66,10 +65,10 @@ at::Tensor glmLowmBf16Moe(at::Tensor const& x, at::Tensor const& routerLogits, a int64_t const hiddenSize = x.size(1); int64_t const numExperts = w3W1.size(0); int64_t const interSize = w2.size(2); - TORCH_CHECK(numExperts == tk::kLowMBf16MoeNumExperts, "glm_lowm_bf16_moe is built for ", - tk::kLowMBf16MoeNumExperts, " experts, got ", numExperts); - TORCH_CHECK(topK == tk::kLowMBf16MoeTopK, "glm_lowm_bf16_moe is built for top-", tk::kLowMBf16MoeTopK, - ", got top-", topK); + TORCH_CHECK(numExperts == tk::kLowMBf16MoeNumExperts, "glm_lowm_bf16_moe is built for ", tk::kLowMBf16MoeNumExperts, + " experts, got ", numExperts); + TORCH_CHECK( + topK == tk::kLowMBf16MoeTopK, "glm_lowm_bf16_moe is built for top-", tk::kLowMBf16MoeTopK, ", got top-", topK); TORCH_CHECK(w3W1.size(1) == 2 * interSize && w3W1.size(2) == hiddenSize, "glm_lowm_bf16_moe: w3_w1 of shape ", w3W1.sizes(), " does not match x ", x.sizes(), " and w2 ", w2.sizes()); TORCH_CHECK(w2.size(0) == numExperts && w2.size(1) == hiddenSize, "glm_lowm_bf16_moe: w2 of shape ", w2.sizes(), diff --git a/cpp/tensorrt_llm/thop/glmMlaKvHalfAppendOp.cpp b/cpp/tensorrt_llm/thop/glmMlaKvHalfAppendOp.cpp index 3726c3f61459..c2d2759a08f4 100644 --- a/cpp/tensorrt_llm/thop/glmMlaKvHalfAppendOp.cpp +++ b/cpp/tensorrt_llm/thop/glmMlaKvHalfAppendOp.cpp @@ -76,8 +76,8 @@ void glmMlaKvHalfAppend(th::Tensor const& latent, th::Tensor const& kvNormWeight "latent rows must be 16-byte aligned (row stride a multiple of 8 elements), got row stride ", latent.stride(0)); int64_t const numTokens = latent.size(0); TORCH_CHECK(numTokens <= kInt32Max, "too many latent rows"); - TORCH_CHECK(seqLenQ > 0 && numTokens % seqLenQ == 0, - "seq_len_q must be positive and divide the latent rows; got ", seqLenQ, " for ", numTokens, " rows"); + TORCH_CHECK(seqLenQ > 0 && numTokens % seqLenQ == 0, "seq_len_q must be positive and divide the latent rows; got ", + seqLenQ, " for ", numTokens, " rows"); int64_t const numGenerations = numTokens / seqLenQ; TORCH_CHECK(kvNormWeight.scalar_type() == th::kBFloat16 && kvNormWeight.is_contiguous() diff --git a/cpp/tensorrt_llm/thop/glmMnnvlHopOp.cpp b/cpp/tensorrt_llm/thop/glmMnnvlHopOp.cpp index dff31244e77e..2504d971622d 100644 --- a/cpp/tensorrt_llm/thop/glmMnnvlHopOp.cpp +++ b/cpp/tensorrt_llm/thop/glmMnnvlHopOp.cpp @@ -80,8 +80,8 @@ tensorrt_llm::runtime::McastDeviceMemory* hopWorkspace( TORCH_CHECK(hopShapeSupported(numTokens, hiddenSize, nRanks), op, ": unsupported shape: ", numTokens, " tokens, hidden ", hiddenSize, ", ", nRanks, " ranks"); int64_t const bufferBytes = comm_buffer.size(1) * comm_buffer.element_size(); - int64_t const hopBytes = glm::mnnvlHopBufferBytes( - static_cast(numTokens), static_cast(hiddenSize), static_cast(nRanks)); + int64_t const hopBytes + = glm::mnnvlHopBufferBytes(static_cast(numTokens), static_cast(hiddenSize), static_cast(nRanks)); TORCH_CHECK(hopBytes <= bufferBytes, op, ": ", hopBytes, " bytes of Lamport data exceed the Lamport buffer of ", bufferBytes, " bytes"); return mcastMem; diff --git a/cpp/tensorrt_llm/thop/glmSparseMlaFlashDecodeOp.cpp b/cpp/tensorrt_llm/thop/glmSparseMlaFlashDecodeOp.cpp index a625400fdc56..019d75820aa9 100644 --- a/cpp/tensorrt_llm/thop/glmSparseMlaFlashDecodeOp.cpp +++ b/cpp/tensorrt_llm/thop/glmSparseMlaFlashDecodeOp.cpp @@ -133,12 +133,13 @@ void glmSparseMlaFlashDecode(th::Tensor const& q, th::Tensor const& kvPool, th:: if (blockTable.has_value()) { auto const& table = blockTable.value(); - TORCH_CHECK(table.is_cuda() && table.device() == device && table.scalar_type() == th::kInt32 - && table.dim() == 2 && table.stride(1) == 1 && table.size(0) >= numSeqs && table.size(1) > 0 - && table.size(1) <= kInt32Max && table.stride(0) >= table.size(1), + TORCH_CHECK(table.is_cuda() && table.device() == device && table.scalar_type() == th::kInt32 && table.dim() == 2 + && table.stride(1) == 1 && table.size(0) >= numSeqs && table.size(1) > 0 && table.size(1) <= kInt32Max + && table.stride(0) >= table.size(1), "block_table must be an int32 CUDA [sequences, max_blocks] tensor with unit column stride"); - TORCH_CHECK(tokensPerBlock > 0 && pageIndexScale > 0 && layerOffset >= 0, "invalid block_table geometry: ", - "tokens_per_block ", tokensPerBlock, ", page_index_scale ", pageIndexScale, ", layer_offset ", layerOffset); + TORCH_CHECK(tokensPerBlock > 0 && pageIndexScale > 0 && layerOffset >= 0, + "invalid block_table geometry: ", "tokens_per_block ", tokensPerBlock, ", page_index_scale ", + pageIndexScale, ", layer_offset ", layerOffset); int64_t const strideFactor = pageIndexScale * tokensPerBlock; int64_t const rowOffset = layerOffset * tokensPerBlock; TORCH_CHECK(strideFactor <= kInt32Max && rowOffset <= kInt32Max, "block_table rows must be int32-addressable"); diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py index 2bd0708e6418..78c66480d1ab 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py @@ -31,9 +31,7 @@ AttentionInputType, AttentionMetadata, ) -from tensorrt_llm._torch.attention.backends.sparse.dsa.module import ( - forward_generation_sparse_attn, -) +from tensorrt_llm._torch.attention.backends.sparse.dsa.module import forward_generation_sparse_attn from tensorrt_llm._torch.attention.backends.sparse.dsa.params import DSABackendForwardArgs from tensorrt_llm._torch.attention.backends.sparse.params import MTPIndexShareMetadata from tensorrt_llm._torch.distributed import AllReduceFusionOp, AllReduceParams @@ -341,12 +339,14 @@ def _check_attention(self, layer) -> None: """Assert the layer's MLA weights and KV cache are the layout the kernels take.""" mla = layer.self_attn assert not mla.apply_rotary_emb and not mla.mapping.has_cp_helix() - assert tuple(mla.kv_a_proj_with_mqa.weight.shape) == ( - FUSED_A_OUT_FEATURES, - self.hidden_size, - ) and mla.kv_a_proj_with_mqa.weight.dtype == torch.bfloat16, ( - f"layer {layer.layer_idx}: kv_a_proj_with_mqa is not the BF16 [q_a; kv_a] weight" - ) + assert ( + tuple(mla.kv_a_proj_with_mqa.weight.shape) + == ( + FUSED_A_OUT_FEATURES, + self.hidden_size, + ) + and mla.kv_a_proj_with_mqa.weight.dtype == torch.bfloat16 + ), f"layer {layer.layer_idx}: kv_a_proj_with_mqa is not the BF16 [q_a; kv_a] weight" assert ( getattr(mla.mqa, "has_fp8_kv_cache", False) and mla.kv_cache_dtype != "fp8_ds_mla" @@ -1095,9 +1095,7 @@ def select_then_q_b(): self.index_topk, ) return output - page_index_scale, layer_offset = manager.get_primary_pool_page_index_params( - local_layer - ) + page_index_scale, layer_offset = manager.get_primary_pool_page_index_params(local_layer) sparse_mla_flash_decode( q_heads, self._kv_pool, @@ -1119,9 +1117,7 @@ def select_then_q_b(): return output # The built-in indexer over the whole context, then the selection as cache rows. - maybe_execute_in_parallel( - q_b, kv_half, mla.ln_events[0], mla.ln_events[1], mla.aux_stream - ) + maybe_execute_in_parallel(q_b, kv_half, mla.ln_events[0], mla.ln_events[1], mla.aux_stream) indexer_intermediates = None if indexer is not None: qr = mla.q_a_layernorm(q_a) @@ -1426,10 +1422,11 @@ def _attention_builtin(layer, position_ids, hidden_states, attn_metadata, all_re indexer_intermediates = None if indexer is not None: indexer_intermediates = [ - t[:num_tokens] - for t in indexer.pre_indexer_proj(qr, hidden_states, position_ids) + t[:num_tokens] for t in indexer.pre_indexer_proj(qr, hidden_states, position_ids) ] - indexer._update_k_cache(indexer_intermediates[1], indexer_intermediates[2], attn_metadata) + indexer._update_k_cache( + indexer_intermediates[1], indexer_intermediates[2], attn_metadata + ) output = hidden_states.new_empty([hidden_states.shape[0], mla.attention_output_hidden_size]) forward_generation_sparse_attn( mla, @@ -1520,9 +1517,7 @@ def __init__( mapping, decode_target, ) -> None: - super().__init__( - spec_config, model_config, use_separate_draft_kv_cache, mapping=mapping - ) + super().__init__(spec_config, model_config, use_separate_draft_kv_cache, mapping=mapping) # The core's per-step lookup of the decode target (a bound method rather than the # core itself, which belongs to the shell and is no submodule of this worker). self._decode_target = decode_target @@ -1776,9 +1771,7 @@ class ModelingV2Glm52Nvfp4Sm107Tp16( """GLM-5.2-NVFP4, sm107, tp16.""" @classmethod - def get_preferred_kv_cache_manager_version( - cls, pretrained_config: Any = None - ) -> Literal["V2"]: + def get_preferred_kv_cache_manager_version(cls, pretrained_config: Any = None) -> Literal["V2"]: """The KV cache manager the context server's built-in implementation prefers.""" return "V2" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/mla_kv_half_append.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/mla_kv_half_append.py index d777072dfc1b..948cc51554d8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/mla_kv_half_append.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/mla_kv_half_append.py @@ -219,8 +219,10 @@ def reference( rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + kv_norm_eps) nope = x * (rstd * kv_norm_weight.float()) - pairs = latent[:, KV_LORA_RANK : KV_LORA_RANK + ROPE_DIM].float().reshape( - num_tokens, ROPE_DIM // 2, 2 + pairs = ( + latent[:, KV_LORA_RANK : KV_LORA_RANK + ROPE_DIM] + .float() + .reshape(num_tokens, ROPE_DIM // 2, 2) ) table = rotary_cos_sin.reshape(-1, ROPE_DIM, 2)[positions.clamp(min=0), : ROPE_DIM // 2] cos, sin = table[..., 0], table[..., 1] @@ -229,11 +231,7 @@ def reference( dim=-1, ).reshape(num_tokens, ROPE_DIM) - rows = ( - torch.cat((nope, roped), dim=-1) - .clamp(-_E4M3_MAX, _E4M3_MAX) - .to(torch.float8_e4m3fn) - ) + rows = torch.cat((nope, roped), dim=-1).clamp(-_E4M3_MAX, _E4M3_MAX).to(torch.float8_e4m3fn) scale = torch.tensor(softmax_scale, dtype=torch.float32, device=device) log2e = torch.tensor(_LOG2E, dtype=torch.float32, device=device) bmm1 = torch.stack((scale, scale * log2e)) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits.py index 84b4f42fe54b..71df8ceeac9e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits.py @@ -101,9 +101,7 @@ def _compile( q_fake = cute.runtime.make_fake_compact_tensor( cutlass.Uint8, (n, head_dim, sym_b), stride_order=(1, 0, 2) ) - w_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Float32, (n, sym_b), stride_order=(0, 1) - ) + w_fake = cute.runtime.make_fake_compact_tensor(cutlass.Float32, (n, sym_b), stride_order=(0, 1)) logits_fake = cute.runtime.make_fake_tensor( cutlass.Float32, (cute.sym_int(), max_ctx), stride=(cute.sym_int64(), 1) ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk_fused.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk_fused.py index 4f7a258f7ebf..89b40efdca3b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk_fused.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk_fused.py @@ -1603,9 +1603,7 @@ def _store_qtail( ) if col0 + m4 == 0: mQScaleOut[token, head] = scale - mWOut[token, head] = _mul_rn( - _mul_rn(sW[row], scale), Float32(qt.weight_scale) - ) + mWOut[token, head] = _mul_rn(_mul_rn(sW[row], scale), Float32(qt.weight_scale)) # ----------------------------------------------------------------------------------- host @@ -1719,9 +1717,7 @@ def _validate_fused(a, b, out, epilogue, gamma, norm_out, kv_lens, cos_sin, q_sc _check(b.shape[0] == k, f"incompatible shapes {tuple(a.shape)} x {tuple(b.shape)}") n = b.shape[1] _check(tuple(out.shape) == (m, n), f"out must be {(m, n)}, got {tuple(out.shape)}") - _check( - 1 <= m <= 32 and k % _CTA_K == 0 and n % 4 == 0, f"unsupported shape M={m} N={n} K={k}" - ) + _check(1 <= m <= 32 and k % _CTA_K == 0 and n % 4 == 0, f"unsupported shape M={m} N={n} K={k}") # a may be a row view of a wider buffer (the raw q_a slice of the fused-A output). _check( a.stride(1) == 1 and a.stride(0) >= k and (a.stride(0) * 2) % 16 == 0, @@ -1729,17 +1725,13 @@ def _validate_fused(a, b, out, epilogue, gamma, norm_out, kv_lens, cos_sin, q_sc ) _check(b.t().is_contiguous(), "b must be column-major (the transpose of a contiguous [N, K])") _check(out.is_contiguous(), "out must be contiguous") - _check( - all(t.data_ptr() % 32 == 0 for t in (a, b, out)), "a, b and out must be 32-byte aligned" - ) + _check(all(t.data_ptr() % 32 == 0 for t in (a, b, out)), "a, b and out must be 32-byte aligned") if epilogue.fp8_rope: _check(out.dtype == _torch.uint8, "fp8_rope writes uint8 (E4M3 bits)") _check( n % epilogue.head_dim == 0, f"N={n} is not a multiple of head_dim {epilogue.head_dim}" ) - _check( - m % epilogue.seq_len == 0, f"M={m} is not a multiple of seq_len {epilogue.seq_len}" - ) + _check(m % epilogue.seq_len == 0, f"M={m} is not a multiple of seq_len {epilogue.seq_len}") for t, name, dtype in ( (kv_lens, "kv_lens", _torch.int32), (cos_sin, "cos_sin", _torch.float32), @@ -1752,9 +1744,7 @@ def _validate_fused(a, b, out, epilogue, gamma, norm_out, kv_lens, cos_sin, q_sc and t.device == a.device, f"fp8_rope needs a contiguous {dtype} {name} on the GEMM's device", ) - _check( - kv_lens.numel() >= m // epilogue.seq_len, "kv_lens has fewer entries than sequences" - ) + _check(kv_lens.numel() >= m // epilogue.seq_len, "kv_lens has fewer entries than sequences") _check(q_scale.numel() >= 1, "q_scale is empty") else: _check(out.dtype == _torch.bfloat16, "out must be BF16 without fp8_rope") @@ -1929,9 +1919,7 @@ def _validate_qtail( ) _check(b.t().is_contiguous(), "b must be column-major (the transpose of a contiguous [N, K])") _check(out.is_contiguous() and out.dtype == _torch.uint8, "out must be contiguous uint8") - _check( - all(t.data_ptr() % 32 == 0 for t in (a, b, out)), "a, b and out must be 32-byte aligned" - ) + _check(all(t.data_ptr() % 32 == 0 for t in (a, b, out)), "a, b and out must be 32-byte aligned") _check( epilogue.row_rms and not epilogue.rms_out and not epilogue.fp8_rope, "the indexer Q tail follows the row_rms epilogue alone", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/indexer_q_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/indexer_q_tail.py index 98dcc783122e..a8c071f53b76 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/indexer_q_tail.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/indexer_q_tail.py @@ -16,12 +16,7 @@ from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl from ._low_m_splitk import operands_overlap -from ._low_m_splitk_fused import ( - IndexerQTail, - SplitKEpilogue, - idx_qtail_plan, - run_splitk_idx_qtail, -) +from ._low_m_splitk_fused import IndexerQTail, SplitKEpilogue, idx_qtail_plan, run_splitk_idx_qtail #: GLM-5.2's DSA indexer (replicated on every rank): 32 heads of 128, RoPE on the first 64 #: columns as GPT-J pairs, UE8M0 scales, weight scale 128 ** -0.5 * 32 ** -0.5; q_lora_rank 2048 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/low_m_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/low_m_gemm.py index 048993cf89c8..a232199870d6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/low_m_gemm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/low_m_gemm.py @@ -14,12 +14,7 @@ ) from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl -from ._low_m_splitk import ( - device_limits, - operands_overlap, - run_splitk_dense, - splitk_default_tactic, -) +from ._low_m_splitk import device_limits, operands_overlap, run_splitk_dense, splitk_default_tactic class GlmLowMGemm(OpWrapper): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/routed_experts.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/routed_experts.py index 5ff359826b89..6637fcbbe700 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/routed_experts.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/routed_experts.py @@ -111,23 +111,17 @@ class GlmRoutedExperts(OpWrapper): Cell( why="the c1 verify step: 8 rows routed by the fused router GEMV, the input joined " "right before FC1, FC2 at K = 128, unfinalized for the post-MoE hop", - spec=dict( - tokens=8, fused_router=True, late_join=True, gemm2_k128=True, **_GLM_TP16 - ), + spec=dict(tokens=8, fused_router=True, late_join=True, gemm2_k128=True, **_GLM_TP16), ), Cell( why="one row, a step without MTP: the fused router's single-token instance and " "FC2 at K = 128 over 8 token tiles", - spec=dict( - tokens=1, fused_router=True, late_join=True, gemm2_k128=True, **_GLM_TP16 - ), + spec=dict(tokens=1, fused_router=True, late_join=True, gemm2_k128=True, **_GLM_TP16), ), Cell( why="32 rows (4 requests x 8), past the fused router: routed from the router " "GEMM's logits, FC2 back on the Gemm2 cubin", - spec=dict( - tokens=32, fused_router=False, late_join=False, gemm2_k128=True, **_GLM_TP16 - ), + spec=dict(tokens=32, fused_router=False, late_join=False, gemm2_k128=True, **_GLM_TP16), ), ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_ehproj_table.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_ehproj_table.py index 537bcf3460a2..2e7ae1fa500c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_ehproj_table.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_ehproj_table.py @@ -266,6 +266,7 @@ def is_valid( dsa_layout, req_idx_per_token, ) + # x is stored row by row while the cluster still reads the token table for the later rows, and the norm # after the last one: an output overlapping an input, or the other output, replaces values the kernel has # yet to read or to return. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py index 5b4913b55107..8d34e6bf535e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py @@ -23,7 +23,9 @@ from tensorrt_llm._torch.modules.linear import TensorParallelMode, load_weight_shard -def normalize_model_config(model_config: ModelConfig[PretrainedConfig]) -> ModelConfig[PretrainedConfig]: +def normalize_model_config( + model_config: ModelConfig[PretrainedConfig], +) -> ModelConfig[PretrainedConfig]: """The two config rewrites the built-in implementation applies before building modules. - Legacy checkpoints name the fused q_a / kv_a projection ``fused_a`` in their diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index fb54c0741a39..da37794ef6a7 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -946,8 +946,7 @@ def _(target_tokens: torch.Tensor, draft_tokens: torch.Tensor, force_counter: Optional[torch.Tensor], force_mode: int, force_base: int, force_max: int, force_frac: float, counter_stride: int, slot_stride: int, pool_size: int): - return draft_tokens.new_empty( - (num_contexts + draft_tokens.shape[0], )) + return draft_tokens.new_empty((num_contexts + draft_tokens.shape[0], )) @torch.library.register_fake("trtllm::glm_draft_sample_glue") def _(logits, exchange_buffer, exchange_flags, group_size, draft_tokens, @@ -956,8 +955,8 @@ def _(logits, exchange_buffer, exchange_flags, group_size, draft_tokens, seq_lens, fill_seq_lens, kv_lens, kv_update, rewind_accepted, draft_len, block_offsets, dsa_layout, req_idx_per_token, slot_mapping_fp8, slot_mapping_scale, gen_kv_indptr, - gen_cached_token_indptr, accepted_tokens, batch_indices, - num_accepted, next_new_tokens): + gen_cached_token_indptr, accepted_tokens, batch_indices, num_accepted, + next_new_tokens): return None @torch.library.register_fake("trtllm::glm_draft_ehproj_table") @@ -1030,12 +1029,13 @@ def _(x, gate_up_weight, down_weight, gate_up_tiling=-1, down_tiling=-1): @torch.library.register_fake("trtllm::glm_fused_a_gemm") def _(hidden_states, weight, l2_prefetch): - return hidden_states.new_empty((hidden_states.shape[0], weight.shape[0])) + return hidden_states.new_empty( + (hidden_states.shape[0], weight.shape[0])) @torch.library.register_fake("trtllm::glm_router_gemm") def _(hidden_states, weight): - return hidden_states.new_empty((hidden_states.shape[0], weight.shape[0]), - dtype=torch.float32) + return hidden_states.new_empty( + (hidden_states.shape[0], weight.shape[0]), dtype=torch.float32) @torch.library.register_fake("trtllm::glm_mnnvl_hop") def _(input, residual, gamma, epsilon, comm_buffer, buffer_flags, diff --git a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py index b973aeb0bcfc..305b3e7912a0 100644 --- a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py @@ -658,45 +658,44 @@ def get_tuning_config(cls, @torch.library.custom_op("trtllm::fp4_block_scale_moe_runner", mutates_args=()) -def fp4_block_scale_moe_runner(routing_logits: Optional[torch.Tensor], - routing_bias: Optional[torch.Tensor], - hidden_states: torch.Tensor, - hidden_states_scale: torch.Tensor, - gemm1_weights: torch.Tensor, - gemm1_weights_scale: torch.Tensor, - gemm1_bias: torch.Tensor, - gemm1_alpha: torch.Tensor, - gemm1_beta: torch.Tensor, - gemm1_clamp_limit: torch.Tensor, - gemm2_weights: torch.Tensor, - gemm2_weights_scale: torch.Tensor, - gemm2_bias: torch.Tensor, - output1_scale_scalar: torch.Tensor, - output1_scale_gate_scalar: torch.Tensor, - output2_scale_scalar: torch.Tensor, - num_experts: int, - top_k: int, - n_group: Optional[int], - topk_group: Optional[int], - intermediate_size: int, - local_expert_offset: int, - local_num_experts: int, - routed_scaling_factor: Optional[float], - routing_method_type: int, - do_finalize: bool, - act_type: int = ActType_TrtllmGen.SwiGlu.value, - topk_weights: Optional[torch.Tensor] = None, - topk_ids: Optional[torch.Tensor] = None, - output: Optional[torch.Tensor] = None, - tune_max_num_tokens: int = 8192, - use_dp: bool = False, - hidden_states_ready_event: Optional[int] = None, - router_hidden_states: Optional[ - torch.Tensor] = None, - router_weight: Optional[torch.Tensor] = None, - router_workspace: Optional[torch.Tensor] = None, - prefer_gemm2_k128: bool = False - ) -> List[torch.Tensor]: +def fp4_block_scale_moe_runner( + routing_logits: Optional[torch.Tensor], + routing_bias: Optional[torch.Tensor], + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm1_bias: torch.Tensor, + gemm1_alpha: torch.Tensor, + gemm1_beta: torch.Tensor, + gemm1_clamp_limit: torch.Tensor, + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + gemm2_bias: torch.Tensor, + output1_scale_scalar: torch.Tensor, + output1_scale_gate_scalar: torch.Tensor, + output2_scale_scalar: torch.Tensor, + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: Optional[float], + routing_method_type: int, + do_finalize: bool, + act_type: int = ActType_TrtllmGen.SwiGlu.value, + topk_weights: Optional[torch.Tensor] = None, + topk_ids: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, + tune_max_num_tokens: int = 8192, + use_dp: bool = False, + hidden_states_ready_event: Optional[int] = None, + router_hidden_states: Optional[torch.Tensor] = None, + router_weight: Optional[torch.Tensor] = None, + router_workspace: Optional[torch.Tensor] = None, + prefer_gemm2_k128: bool = False) -> List[torch.Tensor]: """hidden_states_ready_event: raw cudaEvent_t (torch.cuda.Event.cuda_event) recorded on another stream once hidden_states and hidden_states_scale are written; the op waits on it right before FC1, after launching routing. diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_mla_kv_half_append.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_mla_kv_half_append.py index 07147b5ef0d2..c3ad37ea9d02 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_mla_kv_half_append.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_mla_kv_half_append.py @@ -125,9 +125,7 @@ def generation_step(self, cached: list, seq_len_q: int) -> TrtllmAttentionMetada metadata.num_contexts = 0 metadata.request_ids = request_ids metadata.prompt_lens = cached - metadata.kv_cache_params = KVCacheParams( - use_cache=True, num_cached_tokens_per_seq=cached - ) + metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=cached) metadata.prepare() return metadata diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_sparse_mla_flash_decode.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_sparse_mla_flash_decode.py index e2fe77249109..a255b9e9819b 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_sparse_mla_flash_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_sparse_mla_flash_decode.py @@ -196,8 +196,7 @@ def _build(spec: dict, seed: int) -> _Case: mla_bmm1_scale=torch.stack((scale, scale * torch.tensor(_LOG2E))).cuda(), mla_bmm2_scale=torch.ones(1, dtype=torch.float32, device="cuda"), v_b_proj=( - torch.randn(NUM_HEADS, V_HEAD_DIM, KV_LORA_RANK, generator=gen, device="cuda") - * _VB_STD + torch.randn(NUM_HEADS, V_HEAD_DIM, KV_LORA_RANK, generator=gen, device="cuda") * _VB_STD ).to(torch.bfloat16), output=torch.full( (num_tokens + _SPARE_ROWS, NUM_HEADS * V_HEAD_DIM), @@ -411,9 +410,7 @@ def generation_step(self, cached: list, seq_len_q: int) -> TrtllmAttentionMetada metadata.num_contexts = 0 metadata.request_ids = request_ids metadata.prompt_lens = cached - metadata.kv_cache_params = KVCacheParams( - use_cache=True, num_cached_tokens_per_seq=cached - ) + metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=cached) metadata.prepare() return metadata @@ -478,9 +475,9 @@ def test_reads_the_rows_the_kv_append_writes_in_the_dsa_pool() -> None: fused_a = torch.randn(num_tokens, 2624, generator=gen, device="cuda").to(torch.bfloat16) latent = fused_a[:, LATENT_OFFSET : LATENT_OFFSET + HEAD_DIM_QK] - kv_norm_weight = ( - 1.0 + 0.1 * torch.randn(KV_LORA_RANK, generator=gen, device="cuda") - ).to(torch.bfloat16) + kv_norm_weight = (1.0 + 0.1 * torch.randn(KV_LORA_RANK, generator=gen, device="cuda")).to( + torch.bfloat16 + ) bmm1 = torch.full((2,), float("nan"), dtype=torch.float32, device="cuda") bmm2 = torch.full((1,), float("nan"), dtype=torch.float32, device="cuda") append_args = ( @@ -519,8 +516,7 @@ def test_reads_the_rows_the_kv_append_writes_in_the_dsa_pool() -> None: torch.randn(num_tokens, NUM_HEADS, HEAD_DIM_QK, generator=gen, device="cuda") * _Q_STD ).to(torch.float8_e4m3fn) v_b_proj = ( - torch.randn(NUM_HEADS, V_HEAD_DIM, KV_LORA_RANK, generator=gen, device="cuda") - * _VB_STD + torch.randn(NUM_HEADS, V_HEAD_DIM, KV_LORA_RANK, generator=gen, device="cuda") * _VB_STD ).to(torch.bfloat16) output = torch.empty( num_tokens, NUM_HEADS * V_HEAD_DIM, dtype=torch.bfloat16, device="cuda" diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_moe_finalize_hop_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_moe_finalize_hop_op_matrix.py index 2544435a37d6..a6c3e5478206 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_moe_finalize_hop_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_moe_finalize_hop_op_matrix.py @@ -177,8 +177,7 @@ def check_empty_slots_and_wide_rows() -> None: for rows, pitch, empty in ((8, HIDDEN, True), (8, HIDDEN + 128, False), (3, HIDDEN + 64, True)): seed = 2000 + rows + pitch per_rank = [ - _moe_inputs(rows, seed, r, "top8", pitch=pitch, empty_slots=empty) - for r in range(WORLD) + _moe_inputs(rows, seed, r, "top8", pitch=pitch, empty_slots=empty) for r in range(WORLD) ] residual = _residual(rows, seed) norm, residual_out = _fused(per_rank[RANK], residual, gamma) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_rank_job.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_rank_job.py index 3cc315bed18e..7d3e5d460f08 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_rank_job.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_rank_job.py @@ -56,9 +56,7 @@ def run(entry: str) -> None: package = importlib.util.find_spec("tensorrt_llm") assert package is not None and package.origin is not None, "tensorrt_llm is not importable" package_parent = str(Path(package.origin).resolve().parents[1]) - env["PYTHONPATH"] = os.pathsep.join( - p for p in (package_parent, env.get("PYTHONPATH", "")) if p - ) + env["PYTHONPATH"] = os.pathsep.join(p for p in (package_parent, env.get("PYTHONPATH", "")) if p) timeout = _load_launcher_constants(launcher).DEADLINE_S + _LAUNCHER_GRACE_S completed = subprocess.run( [sys.executable, str(launcher)], diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_indexer_inputs.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_indexer_inputs.py index ce7f6d28e8e4..e1419eafb0eb 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_indexer_inputs.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_indexer_inputs.py @@ -53,9 +53,7 @@ def block_row(positions: torch.Tensor, pool_pages: int, gen: torch.Generator) -> return row -def slot_mappings( - positions: torch.Tensor, row: torch.Tensor -) -> Tuple[torch.Tensor, torch.Tensor]: +def slot_mappings(positions: torch.Tensor, row: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """(fp8 slots, scale slots) int64 flat byte offsets of each position's K row, the metadata's ``_compute_slot_mappings`` for one request.""" pos = positions.long().cpu() diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_indexer_q_tail.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_indexer_q_tail.py index c5c26a0a31da..03675a813a89 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_indexer_q_tail.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_indexer_q_tail.py @@ -51,9 +51,7 @@ def _build(spec: dict, seed: int) -> dict: weight = (torch.randn(heads * head_dim, k, generator=g, device="cuda") * 2.0 * k**-0.5).to( torch.bfloat16 ) - positions = torch.randint( - 0, _POSITIONS, (m,), generator=g, device="cuda", dtype=torch.int32 - ) + positions = torch.randint(0, _POSITIONS, (m,), generator=g, device="cuda", dtype=torch.int32) q_fp8 = torch.full((m, heads, head_dim), _E4M3_NAN, dtype=torch.uint8, device="cuda") return dict( q_a=fused_a[:, :k], diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_lowm_bf16_moe.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_lowm_bf16_moe.py index d8e7de4347f2..9c609339c90c 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_lowm_bf16_moe.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_lowm_bf16_moe.py @@ -13,7 +13,11 @@ from tensorrt_llm._utils import get_sm_version __extra_import_path__ = ["../../../.."] -from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 — needs the path declared above +from _validating import ( # noqa: E402,F401 — needs the path declared above + assert_within_ulp, + certifying, + validating, +) pytestmark = pytest.mark.skipif( get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" @@ -59,9 +63,7 @@ def _build(spec, seed): elif routing == "ties": # Half-integer logits in [-2, 2] and no bias: about 28 experts share the top key, so which 8 of them are # selected, and in which order, is the tie-break. - logits = ( - torch.randint(-4, 5, (1, NUM_EXPERTS), generator=g, device="cuda").float() / 2.0 - ) + logits = torch.randint(-4, 5, (1, NUM_EXPERTS), generator=g, device="cuda").float() / 2.0 bias = torch.zeros(NUM_EXPERTS, device="cuda") else: assert routing == "tiny" diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_routed_experts.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_routed_experts.py index 059591716d8f..57a7f2653d88 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_routed_experts.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_routed_experts.py @@ -123,9 +123,7 @@ def _stack(): gt_p, gt_c, gt_s = _rand_nvfp4(EXPERTS, INTER, HIDDEN, gen) dn_p, dn_c, dn_s = _rand_nvfp4(EXPERTS, HIDDEN, INTER, gen) fc1_perm = _interleave_perm(2 * INTER)[_blk32_perm(2 * INTER)] - w1, s1 = _prepare( - torch.cat([up_p, gt_p], dim=1), torch.cat([up_s, gt_s], dim=1), fc1_perm - ) + w1, s1 = _prepare(torch.cat([up_p, gt_p], dim=1), torch.cat([up_s, gt_s], dim=1), fc1_perm) w2, s2 = _prepare(dn_p, dn_s, _blk32_perm(HIDDEN)) _STACK = dict( gemm1_weights=w1, @@ -464,9 +462,7 @@ def test_gemm2_k128_matches_the_cubin(tokens) -> None: """FC2 at K = 128 writes the Gemm2 cubin's rows, bitwise, and runs where its gate says it does.""" case = _case(tokens, seed=4200 + tokens) - common = dict( - hidden_states=case["xq"], hidden_states_scale=case["xsf"], fused_router=False - ) + common = dict(hidden_states=case["xq"], hidden_states_scale=case["xsf"], fused_router=False) k128 = _call(case, gemm2_k128=True, **common) cubin = _call(case, gemm2_k128=False, **common) _assert_same(k128, cubin, f"{tokens} rows") diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_op_matrix.py index cf79b26d6da0..096da8663092 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_op_matrix.py @@ -349,9 +349,7 @@ def _draft_loop_specs(draft_len: int = 7): table_steps = [] for step in range(draft_len - 1): table_steps.append( - table_op.CELLS[0].spec - if step == 0 - else dict(table_op.CELLS[1].spec, step=step) + table_op.CELLS[0].spec if step == 0 else dict(table_op.CELLS[1].spec, step=step) ) glue_step = dict(glue_op.CELLS[0].spec, step=draft_len - 1) return table_steps, glue_step @@ -441,9 +439,7 @@ def check_guards_refuse() -> None: spec = dict(glue_op.CELLS[0].spec) logits = [_logits(spec, seed, r) for r in range(WORLD)] kwargs, _ = _glue_call(spec, seed, logits, None) - wide = torch.zeros( - spec["rows"], spec["draft_len"] + 2, dtype=torch.int32, device="cuda" - ) + wide = torch.zeros(spec["rows"], spec["draft_len"] + 2, dtype=torch.int32, device="cuda") try: with validating(glue_op): glue_op(**dict(kwargs, next_new_tokens=wide)) @@ -569,9 +565,7 @@ def _spawn_ranks(entry: str, expected_cells: int) -> None: TLLM_DISABLE_MPI="1", ) # The children rendezvous through torch.distributed's env://, not the launching step's PMI namespace. - env = { - k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_")) - } + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} print(f"[launcher] spawning {world_size} ranks for {entry}, rendezvous on :{port}", flush=True) children = [ subprocess.Popen( diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_rank_job.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_rank_job.py index 227701102e53..8adeed47174f 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_rank_job.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_rank_job.py @@ -45,9 +45,7 @@ def run(entry: str, expected_cells: int) -> None: package = importlib.util.find_spec("tensorrt_llm") assert package is not None and package.origin is not None, "tensorrt_llm is not importable" package_parent = str(Path(package.origin).resolve().parents[1]) - env["PYTHONPATH"] = os.pathsep.join( - p for p in (package_parent, env.get("PYTHONPATH", "")) if p - ) + env["PYTHONPATH"] = os.pathsep.join(p for p in (package_parent, env.get("PYTHONPATH", "")) if p) spec = importlib.util.spec_from_file_location(launcher.stem, launcher) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py index 4da00457de51..3e211989b671 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py @@ -13,7 +13,11 @@ from tensorrt_llm._utils import get_sm_version __extra_import_path__ = ["../../../.."] -from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 — needs the path declared above +from _validating import ( # noqa: E402,F401 — needs the path declared above + assert_within_ulp, + certifying, + validating, +) pytestmark = pytest.mark.skipif( get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_token_table.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_token_table.py index f3faf4ffbda7..54b0750306bf 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_token_table.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_token_table.py @@ -34,9 +34,9 @@ def test_table_is_the_token_half_of_eh_proj(tp_size, tp_rank) -> None: """ g = torch.Generator(device="cuda").manual_seed(tp_size * 100 + tp_rank) vocab_shard = 96 - weight = ( - torch.randn(HIDDEN, 2 * HIDDEN, generator=g, device="cuda") * HIDDEN**-0.5 - ).to(torch.bfloat16) + weight = (torch.randn(HIDDEN, 2 * HIDDEN, generator=g, device="cuda") * HIDDEN**-0.5).to( + torch.bfloat16 + ) embed = torch.randn(vocab_shard * tp_size, HIDDEN, generator=g, device="cuda").to( torch.bfloat16 ) From dd8b3581b1e7bba17f0365146cf12c501d312101 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Thu, 15 Oct 2026 20:15:58 -0700 Subject: [PATCH 45/53] [None][fix] GLM-5.2 tp16 decode: read the KV pool from the step's cache manager DecodeTarget kept the first KV cache manager's primary pool. The executor tears that manager down after its memory-profiling pass and allocates the final one, so the final executor's first decode warm-up read freed memory (illegal memory access on every rank). The pool view is memoized per manager on main, so it is now asked for at each attention call. Signed-off-by: Dongfeng Yu --- .../glm_5_2_nvfp4__sm_107__tp16/modeling.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py index 78c66480d1ab..c86ebe3b0e2e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py @@ -314,7 +314,6 @@ def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: self._bmm1_scale = torch.empty(2, dtype=torch.float32, device=device) self._bmm2_scale = torch.empty(1, dtype=torch.float32, device=device) self._unit_scale = torch.ones(1, dtype=torch.float32, device=device) - self._kv_pool = None moe_layers = [layer for layer in layers if not isinstance(layer.mlp, GatedMLP)] # One exchange workspace per router call site: its launches must run in stream order. @@ -1006,8 +1005,8 @@ def _attention( cos_sin = mqa.rotary_cos_sin kv_lens = attn_metadata.kv_lens_cuda_runtime local_layer = mqa.get_local_layer_idx(attn_metadata) - if self._kv_pool is None: - self._kv_pool = manager.get_unique_primary_pool() + # Asked per step: the executor re-creates the cache manager after its memory profiling. + kv_pool = manager.get_unique_primary_pool() reuses_selection = pool_rows is not None or positions is not None own_indexer = sharded and indexer is not None and not reuses_selection @@ -1084,7 +1083,7 @@ def select_then_q_b(): if pool_rows is not None: sparse_mla_flash_decode( q_heads, - self._kv_pool, + kv_pool, pool_rows, kv_lens, self._bmm1_scale, @@ -1098,7 +1097,7 @@ def select_then_q_b(): page_index_scale, layer_offset = manager.get_primary_pool_page_index_params(local_layer) sparse_mla_flash_decode( q_heads, - self._kv_pool, + kv_pool, self._hist_buffers["topk"][:num_tokens] if positions is None else positions, kv_lens, self._bmm1_scale, @@ -1143,7 +1142,7 @@ def select_then_q_b(): ) sparse_mla_flash_decode( q_heads, - self._kv_pool, + kv_pool, indices, kv_lens, self._bmm1_scale, From cb024eeaf2e2d5d0d5eb211877e3e8d1c9716040 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Fri, 16 Oct 2026 01:26:51 -0700 Subject: [PATCH 46/53] [None][feat] GLM-5.2 tp16 draft loop: drop the per-step DSA metadata rebuild The draft steps after the first run no indexer and read only the sequence and KV lengths their kernels advance in place, so the built-in loop's update_for_spec_dec (fused DSA metadata, DeepGEMM schedules, 2-D lengths) after every draft step is dead work in the captured graph. With a short context only the next step's dense selection is refreshed. Signed-off-by: Dongfeng Yu --- .../glm_5_2_nvfp4__sm_107__tp16/modeling.py | 30 +++++++++++-------- 1 file changed, 17 insertions(+), 13 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py index c86ebe3b0e2e..2799f1011c61 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py @@ -670,8 +670,8 @@ def draft_next_step( eh_proj output x. The layer runs no indexer. Its rows attend to ``pool_rows``, draft step 0's selection (``join_selection``, on the first such step, joins the side stream that made it); with a context short enough that the built-in metadata skips the - indexer, to every position up to their own, the selection that metadata lists (rebuilt - for this step by ``update_for_spec_dec``). Returns the MTP layer's output rows. + indexer, to every position up to their own, the dense selection that metadata lists + (refreshed for this step by the loop). Returns the MTP layer's output rows. """ draft = self._draft layer = draft.layer @@ -742,9 +742,8 @@ def draft_ehproj_table( the next step's eh_proj output x and ``input_layernorm(x)``, and its positions: ``positions`` (at ``gather_ids`` on draft step 0) plus one. On draft step 0 ``seq_lens`` becomes 1 and ``kv_lens`` drops the rejected drafts - (``rewind_accepted``); afterwards ``kv_lens`` grows by one. The DSA decode metadata - is left to ``update_for_spec_dec``. Returns ``((normed, x), positions)`` for the next - step. + (``rewind_accepted``); afterwards ``kv_lens`` grows by one. Returns + ``((normed, x), positions)`` for the next step. """ draft = self._draft norm = draft.layer.input_layernorm @@ -1501,10 +1500,9 @@ class MTPDraftWorker(MTPEagleWorker): built-in loop. The loop keeps the built-in loop's attention-metadata handling: the MTP index-share flags, - the switch to one row per request after draft step 0, the per-step - ``update_for_spec_dec`` that rebuilds the DSA decode metadata from the sequence and KV - lengths the draft-step kernels write, and the cleanup. The draft tokens and the - next-new-tokens rows come from those kernels. + the switch to one row per request after draft step 0, and the cleanup. The draft-step + kernels advance the sequence and KV lengths in place, and the draft tokens and the + next-new-tokens rows come from them. """ def __init__( @@ -1683,9 +1681,12 @@ def _draft_on_target( ) -> torch.Tensor: """The draft loop on the decode target's kernels: the draft tokens [batch, draft_len]. - After each step's kernel, which writes the next step's sequence and KV lengths, - ``update_for_spec_dec`` rebuilds the metadata derived from them, as in the built-in - loop; the host side of draft step 0's switch to one row per request comes before it. + Each step's last kernel writes the next step's sequence and KV lengths in place. The + steps after the first run no indexer and read nothing else of the DSA decode metadata, + so the built-in loop's per-step rebuild of it (``update_for_spec_dec``) is left out; + with a context short enough that the metadata skips the indexer, only the next step's + dense selection is refreshed from the new lengths. The host side of draft step 0's + switch to one row per request comes before that. """ draft_len = spec_metadata.runtime_draft_len # The accepted tokens' rows: draft step 0's input ids (prepare_1st_drafter_inputs). @@ -1757,7 +1758,10 @@ def _draft_on_target( kv_lens, rewind_accepted, ) - attn_metadata.update_for_spec_dec() + if not last and attn_metadata.skip_indexer_for_gen_reqs: + attn_metadata.prepare_dense_topk_indices( + attn_metadata.kv_lens_cuda, device=True + ) if i == 0: attn_metadata.use_spec_decoding = False self._drafted_rows = (draft_tokens, next_new_tokens) From 5556e16045fbb5e58b594ef93b131c78377dc2ff Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Fri, 16 Oct 2026 01:26:52 -0700 Subject: [PATCH 47/53] [None][feat] Refresh the kv-length metadata once per step under the overlap correction With speculative decoding and the overlap scheduler, _preprocess_inputs called on_update_kv_lens before the kv_lens correction and again after it. The second call supersedes the first and nothing reads the metadata in between, so the first is skipped when the second runs. Signed-off-by: Dongfeng Yu --- .../_torch/pyexecutor/engine/runners/decoder/runner.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index c74a0a8d9e7e..4aeb996a7d6a 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -2856,8 +2856,16 @@ def _preprocess_inputs( Make some changes to the device inputs and avoid blocking the async data transfer """ attn_meta = inputs.get("attn_metadata") + # The overlap correction below changes kv_lens_cuda and refreshes the metadata after it. + refreshed_after_correction = ( + enable_spec_decode + and not self._config.disable_overlap_scheduler + and attn_meta is not None + and attn_meta.kv_cache_manager is not None + and hasattr(attn_meta, "kv_lens_cuda") + ) # Invalidate per-forward-pass caches so they are recomputed (and captured) on every _forward_step. - if attn_meta is not None: + if attn_meta is not None and not refreshed_after_correction: attn_meta.on_update_kv_lens() if enable_spec_decode and not self._config.disable_overlap_scheduler: From a51b5d37804578bcb632f647a24900b5df37dfd1 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Fri, 16 Oct 2026 02:57:09 -0700 Subject: [PATCH 48/53] [None][feat] GLM-5.2 tp16 draft worker: accept draft tokens in one kernel For a verify step of generation requests only, MTPDraftWorker samples the target tokens as the built-in path does and then runs glm_verify_head_accept, which writes the accepted rows and the strict count with the forced count applied (same warm-up skip, pool draw and counter step as _apply_force_accepted_tokens), instead of the 8-20 small kernels of the built-in chain. Ragged verification and context requests keep that chain. Signed-off-by: Dongfeng Yu --- .../glm_5_2_nvfp4__sm_107__tp16/modeling.py | 59 +++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py index 2799f1011c61..c6cd2d5b73d8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py @@ -52,6 +52,11 @@ ) from tensorrt_llm._torch.speculative import MTPEagleWorker, SpecMetadata from tensorrt_llm._torch.speculative.eagle3 import _reset_mtp_index_share +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, +) from tensorrt_llm._torch.utils import EventType from . import weights as _weights @@ -106,6 +111,7 @@ from .ops.spec.draft_glue_rmsnorm_cat import glm_draft_glue_rmsnorm_cat from .ops.spec.draft_sample_glue import glm_draft_sample_glue from .ops.spec.gather_last_accepted_rows import glm_gather_last_accepted_rows +from .ops.spec.verify_head_accept import glm_verify_head_accept # The deployment this target is written for; route() selects it on the same # facts, and the constructor states them again for a direct construction. @@ -1572,6 +1578,59 @@ def _prepare_next_new_tokens( num_accepted_tokens, ) + def _sample_and_accept_draft_tokens_base( + self, + logits, + draft_tokens, + num_contexts, + batch_size, + spec_metadata, + ): + """Strict acceptance and the forced accepted-token count in one kernel + (``glm_verify_head_accept``) for a step of generation requests only, with the built-in + target-token sampling; ragged verification and context requests take the built-in + chain. The forced count follows ``_apply_force_accepted_tokens``: none during eager + CUDA-graph warm-up, the integer part plus one, and with a fractional part the same pool + draw and counter step.""" + if spec_metadata is None or spec_metadata.is_ragged_verify or num_contexts != 0: + return super()._sample_and_accept_draft_tokens_base( + logits, draft_tokens, num_contexts, batch_size, spec_metadata + ) + if logits.dim() == 1: + logits = logits.unsqueeze(0) + tokens_per_request = draft_tokens.shape[-1] + 1 + target_tokens = self._sample_tokens_for_batch(logits, spec_metadata, 0, batch_size) + mode, base, frac = 0, 0, 0.0 + warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + if self.force_num_accepted_tokens != 0.0 and not warmup: + int_part = int(self.force_num_accepted_tokens) + frac = self.force_num_accepted_tokens - int_part + base = min(int_part + 1, tokens_per_request) + mode = 2 if frac > 0.0 and base < tokens_per_request else 1 + pool = counter = None + if mode == 2: + self._ensure_force_accept_rng_state(logits.device) + pool, counter = self._force_accept_rng_pool, self._force_accept_rng_counter + accepted_tokens = torch.empty( + (batch_size, tokens_per_request), dtype=torch.int, device=logits.device + ) + num_accepted_tokens = glm_verify_head_accept( + target_tokens, + draft_tokens, + 0, + accepted_tokens, + pool, + counter, + mode, + base, + tokens_per_request, + frac, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + ) + return accepted_tokens, num_accepted_tokens + def _forward_linear_draft_loop( self, inputs, From 360bc221f9df5575fd65679fef65789af4d1df97 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Fri, 16 Oct 2026 09:06:35 -0700 Subject: [PATCH 49/53] [None][feat] GLM-5.2 tp16 decode: q_b beside the top-k, early KV gather on shared layers The decode target ran q_b after the indexer and on the regular split-K 8 plan on every layer, so neither stream layout of the ported kernels was in use: - Indexer layers (and MTP draft step 0) queue q_b on the aux stream after the KV half, on the beside_topk plan (18 CTAs), so it runs beside the histogram top-k instead of after it. - Layers that attend to the previous indexer layer's selection run q_b on the early_gather plan (split-K 4), the KV half triggers its dependents at entry, and fused-A also prefetches that selection into L2, so the flash decode gathers the old KV rows while q_b runs. Signed-off-by: Dongfeng Yu --- .../glm_5_2_nvfp4__sm_107__tp16/modeling.py | 55 +++++++++++++------ 1 file changed, 38 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py index c6cd2d5b73d8..9f484bf95dee 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py @@ -1014,17 +1014,24 @@ def _attention( kv_pool = manager.get_unique_primary_pool() reuses_selection = pool_rows is not None or positions is not None own_indexer = sharded and indexer is not None and not reuses_selection + # A layer that attends to the selection of the indexer layer before it gathers its old KV + # rows before the flash decode's wait. + early_gather = sharded and indexer is None and not reuses_selection weight = ( self._fused_a_indexer_weight[layer.layer_idx] if own_indexer else mla.kv_a_proj_with_mqa.weight ) - l2_prefetch = ( - self._draft.l2_after_fused_a - if reuses_selection - else self._l2_after_fused_a[layer.layer_idx] - ) + if reuses_selection: + l2_prefetch = self._draft.l2_after_fused_a + elif early_gather: + # Also the selection the flash decode reads first. + l2_prefetch = self._l2_after_fused_a[layer.layer_idx] + [ + self._hist_buffers["topk"][:num_tokens] + ] + else: + l2_prefetch = self._l2_after_fused_a[layer.layer_idx] kv_a = fused_a_gemm(hidden_states, weight, l2_prefetch) kv_a = kv_a[:num_tokens] q_a = kv_a[:, : mla.q_lora_rank] @@ -1033,7 +1040,7 @@ def _attention( (num_tokens, mla.num_heads_tp * head_dim), dtype=torch.uint8 ) - def q_b(): + def q_b(plan="regular"): q_b_fp8_rope( q_a, self._q_b_absorbed[layer.layer_idx], @@ -1045,10 +1052,10 @@ def q_b(): head_dim, mla.qk_rope_head_dim, seq_len_q, - "regular", + plan, ) - def kv_half(): + def kv_half(trigger_at_entry=False): mla_kv_half_append( kv_a[:, mla.q_lora_rank : latent_end], mla.kv_a_layernorm.weight, @@ -1066,6 +1073,7 @@ def kv_half(): self._bmm1_scale, self._bmm2_scale, mla.softmax_scale, + trigger_at_entry, ) output = hidden_states.new_empty( @@ -1073,18 +1081,32 @@ def kv_half(): ) q_heads = quant_q.view(num_tokens, mla.num_heads_tp, head_dim) if sharded: + if own_indexer: - def select_then_q_b(): - if own_indexer: + def select(): self._select_rows( indexer, kv_a, latent_end, q_a, position_ids, attn_metadata, num_tokens, mla ) - q_b() - # The KV half on the aux stream beside the indexer and q_b. - maybe_execute_in_parallel( - select_then_q_b, kv_half, mla.ln_events[0], mla.ln_events[1], mla.aux_stream - ) + def kv_half_then_q_b(): + kv_half() + # Queued after the top-k: 18 CTAs leave the SMs to the indexer beside it. + q_b("beside_topk") + + # The indexer on this stream; the KV half and q_b beside it on the aux stream. + maybe_execute_in_parallel( + select, kv_half_then_q_b, mla.ln_events[0], mla.ln_events[1], mla.aux_stream + ) + else: + # With the early gather, q_b at split-K 4 leaves room for the flash decode, which + # the KV half lets launch at its entry. + maybe_execute_in_parallel( + lambda: q_b("early_gather" if early_gather else "regular"), + lambda: kv_half(early_gather), + mla.ln_events[0], + mla.ln_events[1], + mla.aux_stream, + ) if pool_rows is not None: sparse_mla_flash_decode( q_heads, @@ -1115,8 +1137,7 @@ def select_then_q_b(): manager.tokens_per_block, page_index_scale, layer_offset, - # The selection of an earlier indexer layer: its rows can load before the wait. - positions is None and not own_indexer, + early_gather, ) return output From 23791cb00885dac1fde80522913df6e5c31d27b2 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Fri, 16 Oct 2026 09:47:02 -0700 Subject: [PATCH 50/53] [None][feat] Leave the DSA decode metadata of a speculative decode graph to its graph head DSAtrtllmAttentionMetadata gets the engine-only defer_cuda_graph_decode_outputs() scope CSA2 already has. Inside it, prepare() of a generation-only CUDA-graph batch on a native DSA cache (no ragged verification, no helix, no separate draft KV cache) skips the outputs that on_update_kv_lens() rebuilds before any consumer: - the generation cached-token and KV indptr uploads; - req_idx_per_token; - the host dense TopK indices of short contexts; - the indexer slot mappings; - the 2-D KV-length fill, the DeepGEMM schedules and their copies. The host scalars (max_gen_seq_len, the indexer-skip flags) are still set, and the INDEX_KEY offsets, the block table and the speculative expansion buffers are still prepared. Entry 0 of both generation indptrs is zeroed once at allocation, so the unfused rebuild, which writes only the entries after it, leaves the same buffer as the upload did. The INDEX_KEY clamp now covers only the rows past the batch: the batch rows were just copied from pool slots that the cache manager clamps to be non-negative, so the clamp never changed them. DecoderRunner._cuda_graph_metadata_prepare_scope now also admits a step of linear one-model speculative decoding that verifies drafts, without a separate draft KV cache manager and without ragged verification. This is correct because _forward_step starts with _preprocess_inputs, which is captured with the graph and also runs in the eager warm-up forwards: - with the overlap scheduler it adds the previous step's acceptance correction to position_ids and kv_lens_cuda and then calls on_update_kv_lens(), which recomputes every skipped output from the corrected lengths, the seq_lens and the INDEX_KEY offsets on the device, before the first layer reads them; - without the overlap scheduler it calls on_update_kv_lens() first, from the prepared lengths; - the draft loop changes these buffers only on the device after the target layers, and the next step's graph head rebuilds them again. Before this change prepare() computed the same outputs from the uncorrected host lengths and on_update_kv_lens() overwrote them before any read, so the values the model reads are unchanged. A separate draft KV cache swaps and refreshes DSA buffers outside the graph, and ragged verification builds per-row layouts on the host; both keep the full preparation, as do direct prepare() callers and DeepSeek-V4, whose metadata subclass has its own prepare() and cache layout. A speculative step without drafts also keeps it, so CSA2, which defers only single-token rows, sees no change. Signed-off-by: Dongfeng Yu --- .../attention/backends/sparse/dsa/indexer.py | 13 ++- .../attention/backends/sparse/dsa/metadata.py | 85 +++++++++++++++++-- .../engine/runners/decoder/runner.py | 28 +++++- 3 files changed, 110 insertions(+), 16 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py index abe993e31438..bfd46a7f78e7 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py @@ -1248,7 +1248,9 @@ def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata): ) @staticmethod - def prepare(metadata: DSAtrtllmAttentionMetadata) -> None: + def prepare( + metadata: DSAtrtllmAttentionMetadata, *, defer_decode_outputs: bool = False + ) -> None: """ Prepare indexer for the forward pass. This should be called during metadata.prepare() stage. @@ -1256,6 +1258,9 @@ def prepare(metadata: DSAtrtllmAttentionMetadata) -> None: - Computes slot_mapping for KV cache updates - Prepares schedule_metadata for fp8_paged_mqa_logits - Stores generation request IDs for decode phase + + ``defer_decode_outputs`` leaves the slot mappings and schedules of a + generation-only batch to ``on_update_kv_lens``. """ # Skip indexer preparation if the kv_cache_manager doesn't have index_head_dim. @@ -1275,9 +1280,11 @@ def prepare(metadata: DSAtrtllmAttentionMetadata) -> None: ) # Store compressed KV token count for context requests metadata.num_ctx_kv_tokens = indexer_params.new_kv_tokens[:num_contexts].sum().item() + defer_decode_outputs = defer_decode_outputs and num_contexts == 0 # Prepare for update_k_cache - Indexer.prepare_for_update_k_cache(metadata, indexer_params) + if not defer_decode_outputs: + Indexer.prepare_for_update_k_cache(metadata, indexer_params) # Prepare for prefill phase if there are context requests if num_contexts > 0: @@ -1305,7 +1312,7 @@ def prepare(metadata: DSAtrtllmAttentionMetadata) -> None: Indexer.prepare_for_chunked_prefill(metadata, indexer_params) # Prepare for decode phase if there are generation requests - if num_generations > 0: + if num_generations > 0 and not defer_decode_outputs: # Prepare schedule metadata for fp8_paged_mqa_logits # This is a preprocessing step that computes scheduling information for the kernel Indexer.prepare_scheduler_metadata(metadata) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py index c6af4efd7dba..ae4188c7b20b 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py @@ -6,6 +6,7 @@ import functools import os +from contextlib import contextmanager from dataclasses import dataclass, field from typing import TYPE_CHECKING, ClassVar, FrozenSet, List, Optional, Tuple @@ -279,6 +280,8 @@ def __init__(self, *args, **kwargs): self.num_ctx_mla_kv_tokens = 0 self.nvfp4_mla_context_fp8_scratch = None self._radix_rows_per_sequence = 1 + # Set only inside defer_cuda_graph_decode_outputs(). + self._defer_decode_outputs = False super().__init__(*args, **kwargs) sparse_metadata_params = self.sparse_metadata_params if not isinstance(sparse_metadata_params, DSAMetadataParams): @@ -357,10 +360,46 @@ def __post_init__(self): self.create_buffers_for_indexer(capture_graph=capture_graph) self._create_nvfp4_mla_generation_buffers(capture_graph=capture_graph) + @contextmanager + def defer_cuda_graph_decode_outputs(self): + """Engine-only scope for a selected decode graph's preparation. + + The engine must run ``on_update_kv_lens`` before any model consumer, + including warmup and capture. That call rebuilds the generation + indptrs, ``req_idx_per_token``, the indexer slot mappings, the DeepGEMM + schedules and the dense TopK indices from the KV lengths the forward + uses, so ``prepare`` leaves them to it for a generation-only batch. + Direct ``prepare`` callers retain fully initialized outputs. + """ + previous = getattr(self, "_defer_decode_outputs", False) + self._defer_decode_outputs = True + try: + yield + finally: + self._defer_decode_outputs = previous + + def _defers_decode_outputs(self) -> bool: + """Whether this ``prepare`` leaves the decode outputs to ``on_update_kv_lens``. + + Ragged verification, helix, a separate draft KV cache and non-native + DSA cache layouts keep the complete preparation. + """ + return ( + getattr(self, "_defer_decode_outputs", False) + and self.is_cuda_graph + and self.num_contexts == 0 + and self.num_generations > 0 + and not self.enable_ragged_verification + and not self.enable_helix + and self.draft_kv_cache_manager is None + and is_dsa_cache_manager(self.kv_cache_manager) + ) + def prepare(self): self._validate_ragged_verification() super().prepare() self._invalidate_pool_view_cache() + defer_decode_outputs = self._defers_decode_outputs() # Get kv lengths assert self.kv_cache_params.use_cache is True, "DSA requires use_cache to be True" @@ -382,13 +421,16 @@ def prepare(self): kv_lens = self._prepare_device_window_kv_lens(kv_lens) # For mla_rope_append_paged_kv_assign_q - self.prepare_for_mla_rope_append(cached_token_lens, kv_lens) + self.prepare_for_mla_rope_append( + cached_token_lens, kv_lens, generation_indptrs=not defer_decode_outputs + ) # Prepare to support skip indexer - self.prepare_for_skip_indexer(kv_lens) + self.prepare_for_skip_indexer(kv_lens, dense_topk_indices=not defer_decode_outputs) # For indices conversion - self.prepare_for_indices_conversion() + if not defer_decode_outputs: + self.prepare_for_indices_conversion() # For indexer k cache self.prepare_for_indexer_k_cache() @@ -397,7 +439,7 @@ def prepare(self): self.prepare_for_spec_decode(kv_lens) # Prepare metadata for indexer - Indexer.prepare(metadata=self) + Indexer.prepare(metadata=self, defer_decode_outputs=defer_decode_outputs) def _prepare_device_window_kv_lens(self, kv_lens: torch.Tensor) -> torch.Tensor: """Bound host-side KV extents until device-selected windows are installed.""" @@ -886,6 +928,10 @@ def create_buffers_for_mla_rope_append(self, capture_graph=False): device="cpu", pin_memory=prefer_pinned(), ) + # Entry 0 of an exclusive prefix sum is always 0, and the unfused + # rebuild in on_update_kv_lens() writes only the entries after it. + self.gen_cached_token_indptr[:1].zero_() + self.gen_kv_indptr[:1].zero_() def _create_nvfp4_mla_generation_buffers(self, capture_graph=False): """Allocate stable NVFP4 gather outputs for all MTP generation rows.""" @@ -1457,9 +1503,12 @@ def _update_indexer_k_cache_block_offsets(self) -> torch.Tensor: self.indexer_k_cache_block_offsets[: self.num_seqs].copy_( self.host_indexer_k_cache_block_offsets[: self.num_seqs], non_blocking=True ) - # Sanitize graph-padding entries that may be stale after cache + # The batch rows hold non-negative pool slots. Sanitize the + # graph-padding rows after them, which may be stale after cache # eviction or host-cache onboarding. - self.indexer_k_cache_block_offsets.clamp_(min=0) + padding_rows = self.indexer_k_cache_block_offsets[self.num_seqs :] + if padding_rows.numel() > 0: + padding_rows.clamp_(min=0) return pool_indices def set_skip_topk(self, skip: bool) -> None: @@ -2028,7 +2077,12 @@ def prepare_for_indexer_k_cache(self): host_block_table, non_blocking=True ) - def prepare_for_skip_indexer(self, kv_lens: torch.Tensor): + def prepare_for_skip_indexer(self, kv_lens: torch.Tensor, dense_topk_indices: bool = True): + """Decide which requests skip the indexer. + + ``dense_topk_indices=False`` leaves their dense TopK indices to + ``on_update_kv_lens``. + """ num_extra_kv_tokens = self.kv_cache_params.num_extra_kv_tokens if self.num_contexts > 0 and self.enable_indexer_skip: # Minus the number of extra KV tokens because when using one-model MTP, the @@ -2049,9 +2103,20 @@ def prepare_for_skip_indexer(self, kv_lens: torch.Tensor): ) else: self.skip_indexer_for_gen_reqs = False - self.prepare_dense_topk_indices(kv_lens) + if dense_topk_indices: + self.prepare_dense_topk_indices(kv_lens) - def prepare_for_mla_rope_append(self, cached_token_lens: torch.Tensor, kv_lens: torch.Tensor): + def prepare_for_mla_rope_append( + self, + cached_token_lens: torch.Tensor, + kv_lens: torch.Tensor, + generation_indptrs: bool = True, + ): + """Set the MLA RoPE-append lengths and indptrs. + + ``generation_indptrs=False`` leaves the generation indptrs to + ``on_update_kv_lens``. + """ if self.num_contexts > 0: self.num_ctx_cached_tokens = cached_token_lens[: self.num_contexts].sum().item() self.num_ctx_mla_kv_tokens = kv_lens[: self.num_contexts].sum().item() @@ -2085,6 +2150,8 @@ def prepare_for_mla_rope_append(self, cached_token_lens: torch.Tensor, kv_lens: if self.num_generations > 0: self.max_gen_seq_len = self.seq_lens[self.num_contexts : self.num_seqs].max().item() + if not generation_indptrs: + return # generation cached token indptr torch.cumsum( cached_token_lens[self.num_contexts : self.num_seqs], diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 4aeb996a7d6a..5b545b3e186b 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -5254,8 +5254,13 @@ def _cuda_graph_metadata_prepare_scope( """Allow derived outputs only when this forward owns a decode graph. Both capture/warmup and replay execute _preprocess_inputs before model - consumers. Non-speculative overlap has the same ordering; speculative - acceptance corrections and draft-manager swaps keep full preparation. + consumers, and it runs on_update_kv_lens, which rebuilds the derived + outputs. Non-speculative overlap has the same ordering. When a linear + one-model speculative step verifies drafts, that call follows the + overlap acceptance correction of kv_lens_cuda, so the outputs come + from the corrected lengths. Draft-manager swaps and ragged + verification rebuild metadata outside the graph and keep full + preparation, as does a speculative step without drafts. """ defer = getattr(attn_metadata, "defer_cuda_graph_decode_outputs", None) if ( @@ -5263,11 +5268,14 @@ def _cuda_graph_metadata_prepare_scope( and isinstance(key, KeyType) and attn_metadata.is_cuda_graph is True and key.num_contexts == 0 - and key.draft_len == 0 + and ( + key.draft_len > 0 and self._linear_one_model_spec_decode(attn_metadata) + if enable_spec_decode + else key.draft_len == 0 + ) and scheduled_requests.num_context_requests == 0 and scheduled_requests.num_generation_requests > 0 and not promoted_context_request_ids - and not enable_spec_decode and not self.use_beam_search and self._context_graph_promotion_supported and self.mapping.cp_size == 1 @@ -5275,6 +5283,18 @@ def _cuda_graph_metadata_prepare_scope( return defer() return contextlib.nullcontext() + def _linear_one_model_spec_decode(self, attn_metadata: AttentionMetadata) -> bool: + """Whether speculation verifies a linear draft in the target model's + own forward, with no separate draft KV cache and no ragged windows.""" + spec_config = self._config.spec_config + return ( + spec_config is not None + and spec_config.spec_dec_mode.use_one_engine() + and spec_config.is_linear_tree + and getattr(attn_metadata, "draft_kv_cache_manager", None) is None + and not getattr(attn_metadata, "enable_ragged_verification", False) + ) + def _get_remote_tail_decode_rank_tokens( self, attn_metadata: AttentionMetadata, can_run_graph: bool ) -> list[int]: From 606312b86e04bf4a9ffd78ea4740d0e78af21bba Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Fri, 16 Oct 2026 09:54:20 -0700 Subject: [PATCH 51/53] [None][feat] Skip unchanged or unread uploads in the speculative step staging Between two decode graph replays, the overlap staging of a speculative step and the sampler issue several copies whose payload is constant for a request, or that nothing reads. This removes them without changing any buffer the step reads. - num_accepted_draft_tokens: nothing in tensorrt_llm reads SpecMetadata.num_accepted_draft_tokens, so the per-request list, its H2D copy and its device buffer are dropped. - previous_pos_id_offsets_cuda / previous_kv_lens_offsets_cuda: the runner zeroed both buffers whole before the previous-batch gathers rewrote their rows. It now tracks the row span each buffer may be nonzero in and clears only that span's rows outside the rows the gathers rewrite next, so the contents after staging are the same and a steady-state step issues no clearing kernel. - host_staging.copy_if_changed() skips an upload whose destination storage still holds the same payload from the last copy made through it (same elements, equal values). The record is keyed by the storage's base tensor, so views of a buffer shared between owners share it. It is used for the previous-batch slot indices, the previous-batch position indices, the gather ids, the Eagle3 batch_indices_cuda and the SpecSampler store slots, which now live in a persistent buffer instead of a new pageable-staged tensor per step. The one other writer of those buffers, the encoder-decoder input fast path, calls forget_copy() after its write. At batch size 1 this removes one H2D for the accepted counts, two zeroing kernels and five constant H2D copies per step. Signed-off-by: Dongfeng Yu --- tensorrt_llm/_torch/host_staging.py | 64 ++++++++++- .../engine/runners/decoder/runner.py | 102 +++++++++++------- .../engine/runners/encoder_decoder.py | 3 + tensorrt_llm/_torch/speculative/eagle3.py | 8 +- .../_torch/speculative/spec_sampler_base.py | 9 +- 5 files changed, 138 insertions(+), 48 deletions(-) diff --git a/tensorrt_llm/_torch/host_staging.py b/tensorrt_llm/_torch/host_staging.py index 38d25db83c7e..43e01119fb03 100644 --- a/tensorrt_llm/_torch/host_staging.py +++ b/tensorrt_llm/_torch/host_staging.py @@ -1,13 +1,75 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Bounded pinned staging for asynchronous host-to-device metadata uploads.""" +"""Asynchronous host-to-device metadata uploads. + +Bounded pinned staging, and skipping an upload whose payload the destination +already holds. +""" import torch +from torch.utils.weak import WeakIdKeyDictionary from tensorrt_llm._utils import maybe_pin_memory, prefer_pinned _STAGING_RING = 2 +# The last copy copy_if_changed() made into each storage, keyed by the +# storage's base tensor: views of one buffer share an entry, and the entry +# goes away with the buffer. +_last_copies = WeakIdKeyDictionary() + + +def _base(tensor: torch.Tensor) -> torch.Tensor: + return tensor if tensor._base is None else tensor._base + + +def _same_payload(recorded, current) -> bool: + if isinstance(recorded, torch.Tensor) or isinstance(current, torch.Tensor): + return ( + isinstance(recorded, torch.Tensor) + and isinstance(current, torch.Tensor) + and recorded.dtype == current.dtype + and recorded.shape == current.shape + and torch.equal(recorded, current) + ) + return recorded == current + + +def copy_if_changed(destination: torch.Tensor, values, *, token=None) -> bool: + """Copy ``values`` into ``destination`` unless it already holds them. + + ``values`` is a host tensor or a sequence of integers. ``destination`` + already holds them when the last copy made here into its storage wrote + the same elements with an equal payload. ``token``, when given, is + compared in place of the payload; it must change whenever the payload + does. A destination filtered here must not be written any other way + without a ``forget_copy`` call. Returns whether a copy was issued. + """ + if destination.numel() == 0: + return False + payload = values if isinstance(values, torch.Tensor) else tuple(values) + key = payload if token is None else token + region = (destination.data_ptr(), destination.dtype, destination.shape, destination.stride()) + base = _base(destination) + last = _last_copies.get(base) + if last is not None and last[0] == region and _same_payload(last[1], key): + return False + if not isinstance(payload, torch.Tensor): + payload = torch.tensor( + payload, dtype=destination.dtype, pin_memory=destination.is_cuda and prefer_pinned() + ) + destination.copy_(payload, non_blocking=destination.is_cuda) + _last_copies[base] = (region, key.clone() if isinstance(key, torch.Tensor) else key) + return True + + +def forget_copy(destination: torch.Tensor) -> None: + """Drop what ``copy_if_changed`` recorded for ``destination``'s storage. + + Call it after writing that storage any other way. + """ + _last_copies.pop(_base(destination), None) + def copy_host_to_device(store: dict, key, destination: torch.Tensor, value: torch.Tensor) -> None: """Asynchronously copy a host tensor into ``destination`` through pinned staging. diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 5b545b3e186b..0357356d4d59 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -22,6 +22,7 @@ from tensorrt_llm._torch.compilation.piecewise_optimizer import PiecewiseRunner from tensorrt_llm._torch.compilation.utils import capture_piecewise_cuda_graph from tensorrt_llm._torch.distributed import Distributed +from tensorrt_llm._torch.host_staging import copy_if_changed from tensorrt_llm._torch.memory_buffer_utils import clear_memory_buffers, with_shared_pool from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.models.modeling_multimodal_mixin import _build_request_multimodal_input @@ -151,6 +152,28 @@ def _zero_fill_row_spans(buffer: torch.Tensor, spans: list[tuple[int, int]]) -> buffer[merged_begin:merged_end].fill_(0) +def _zero_rows_outside( + buffer: torch.Tensor, nonzero_rows: tuple[int, int], kept_rows: tuple[int, int] +) -> tuple[int, int]: + """Zero ``buffer`` everywhere except the half-open row span ``kept_rows``. + + ``buffer`` is zero outside the half-open span ``nonzero_rows``, so only + that span's rows outside ``kept_rows`` are cleared. Returns the span + outside which the buffer is zero once ``kept_rows`` is rewritten. + """ + begin, end = nonzero_rows + kept_begin, kept_end = kept_rows + if kept_begin >= kept_end: + if begin < end: + buffer[begin:end].zero_() + return (0, 0) + if begin < min(end, kept_begin): + buffer[begin : min(end, kept_begin)].zero_() + if max(begin, kept_end) < end: + buffer[max(begin, kept_end) : end].zero_() + return kept_rows + + def resolve_mamba_metadata_cls(model: torch.nn.Module) -> type[Mamba2Metadata]: """Resolve the model-specific Mamba metadata class with a default.""" return getattr(model, "mamba_metadata_cls", None) or Mamba2Metadata @@ -470,9 +493,6 @@ def _allocate_decoder_buffers(self) -> None: self.gather_ids_cuda = torch.empty( (self._config.max_num_tokens,), dtype=torch.int, device="cuda" ) - self.num_accepted_draft_tokens_cuda = torch.empty( - (self._config.max_batch_size,), dtype=torch.int, device="cuda" - ) self.previous_pos_indices_cuda = torch.empty( (self._config.max_num_tokens,), dtype=torch.int, device="cuda" ) @@ -482,6 +502,9 @@ def _allocate_decoder_buffers(self) -> None: self.previous_kv_lens_offsets_cuda = torch.zeros( (self._config.max_batch_size,), dtype=torch.int, device="cuda" ) + # Half-open row spans outside which the two offset buffers are zero. + self._previous_pos_id_offsets_rows = (0, 0) + self._previous_kv_lens_offsets_rows = (0, 0) self.previous_batch_indices_cuda = torch.empty( (self._config.max_num_tokens,), dtype=torch.int, device="cuda" ) @@ -3327,6 +3350,22 @@ def _build_multimodal_data_for_llm( return request.py_multimodal_data return self._mm_item_scheduler.build_multimodal_data_for_llm(request, runtime) + def _zero_previous_offsets_outside( + self, pos_id_rows: tuple[int, int], kv_lens_rows: tuple[int, int] + ) -> None: + """Zero the overlap position-id and KV-length offsets outside the + half-open row spans the caller rewrites next. + + Only these writes change the two buffers, so rows outside the spans + rewritten last are already zero and are not cleared again. + """ + self._previous_pos_id_offsets_rows = _zero_rows_outside( + self.previous_pos_id_offsets_cuda, self._previous_pos_id_offsets_rows, pos_id_rows + ) + self._previous_kv_lens_offsets_rows = _zero_rows_outside( + self.previous_kv_lens_offsets_cuda, self._previous_kv_lens_offsets_rows, kv_lens_rows + ) + def _prepare_tp_inputs( self, scheduled_requests: ScheduledRequests, @@ -3429,7 +3468,6 @@ def _prepare_tp_inputs( mrope_dummy_seq_slot = get_mrope_dummy_seq_slot( self._config.max_num_tokens, self.mapping.pp_size ) - num_accepted_draft_tokens = [] # per request extra_inputs = self._new_extra_inputs() context_prompt_lookahead = None @@ -3466,7 +3504,6 @@ def _prepare_tp_inputs( gather_ids.append(len(input_ids) - 1) sequence_lengths.append(len(prompt_tokens)) - num_accepted_draft_tokens.append(len(prompt_tokens) - 1) prompt_lengths.append(len(prompt_tokens)) past_seen_token_num = begin_compute num_cached_tokens_per_seq.append(past_seen_token_num - request.py_num_compressed_tokens) @@ -3729,7 +3766,6 @@ def _helix_pack_extend(request, group: int) -> int: prompt_lengths.append(request.py_prompt_len) sequence_lengths.append(1 + num_draft_tokens) - num_accepted_draft_tokens.append(num_draft_tokens) gather_ids.extend( list(range(len(position_ids), len(position_ids) + 1 + num_draft_tokens)) ) @@ -3769,7 +3805,6 @@ def _helix_pack_extend(request, group: int) -> int: request.py_batch_idx = request.py_seq_slot sequence_lengths.append(runtime_tokens_per_gen_step) - num_accepted_draft_tokens.append(request.py_num_accepted_draft_tokens) past_seen_token_num = request.max_beam_num_tokens - 1 draft_lens.append(runtime_draft_token_buffer_width) @@ -3834,7 +3869,6 @@ def _helix_pack_extend(request, group: int) -> int: - 1 - (self._config.original_max_draft_len - request.py_num_accepted_draft_tokens) ) - num_accepted_draft_tokens.append(request.py_num_accepted_draft_tokens) sequence_lengths.append(1 + self._config.original_max_draft_len) prompt_lengths.append(request.py_prompt_len) @@ -3897,10 +3931,9 @@ def _helix_pack_extend(request, group: int) -> int: ) # Pre-extend constant-value lists to avoid per-request append - # overhead (saves ~3 append calls per request). + # overhead (saves ~2 append calls per request). draft_lens.extend([0] * (_n_gen * beam_width)) sequence_lengths.extend([1] * (_n_gen * beam_width)) - num_accepted_draft_tokens.extend([0] * (_n_gen * beam_width)) for request in generation_requests: request_ids.append(request.py_request_id) @@ -4066,11 +4099,8 @@ def _helix_pack_extend(request, group: int) -> int: previous_batch_len = len(previous_batch_indices) def previous_seq_slots_device(): - previous_batch_indices_host = torch.tensor( - previous_batch_indices, dtype=torch.int, pin_memory=prefer_pinned() - ) previous_slots = self.previous_batch_indices_cuda[:previous_batch_len] - previous_slots.copy_(previous_batch_indices_host, non_blocking=True) + copy_if_changed(previous_slots, previous_batch_indices) return previous_slots if scheduled_requests.is_attention_dp_phase_idle: @@ -4091,19 +4121,24 @@ def previous_seq_slots_device(): if num_draft_tokens > 0: draft_tokens = torch.tensor(draft_tokens, dtype=torch.int, pin_memory=prefer_pinned()) self.draft_tokens_cuda[: len(draft_tokens)].copy_(draft_tokens, non_blocking=True) - if self._config.is_spec_decode and len(num_accepted_draft_tokens) > 0: - num_accepted_draft_tokens = torch.tensor( - num_accepted_draft_tokens, dtype=torch.int, pin_memory=prefer_pinned() - ) - self.num_accepted_draft_tokens_cuda[: len(num_accepted_draft_tokens)].copy_( - num_accepted_draft_tokens, non_blocking=True - ) if next_draft_tokens_device is not None: - # Initialize these two values to zeros - self.previous_pos_id_offsets_cuda *= 0 - self.previous_kv_lens_offsets_cuda *= 0 runtime_tokens_per_gen_step = self.get_runtime_tokens_per_gen_step(runtime_draft_len) runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 + # Initialize these two values to zeros, except the rows of the + # previous-batch requests, which are written below. + num_extend_reqeust_wo_dummy = len(extend_requests) - len(extend_dummy_requests) + previous_request_rows = ( + (num_extend_reqeust_wo_dummy - previous_batch_len, num_extend_reqeust_wo_dummy) + if previous_batch_len > 0 + else (0, 0) + ) + self._zero_previous_offsets_outside( + ( + previous_request_rows[0] * runtime_tokens_per_gen_step, + previous_request_rows[1] * runtime_tokens_per_gen_step, + ), + previous_request_rows, + ) if previous_batch_len > 0: previous_slots = previous_seq_slots_device() @@ -4129,11 +4164,8 @@ def previous_seq_slots_device(): ) # prepare data for the preprocess inputs kv_len_offsets_device = new_tokens_lens_device - runtime_tokens_per_gen_step - previous_pos_indices_host = torch.tensor( - previous_pos_indices, dtype=torch.int, pin_memory=prefer_pinned() - ) - self.previous_pos_indices_cuda[0:previous_batch_tokens].copy_( - previous_pos_indices_host, non_blocking=True + copy_if_changed( + self.previous_pos_indices_cuda[0:previous_batch_tokens], previous_pos_indices ) # The order of requests in a batch: [context requests, generation requests] @@ -4156,7 +4188,6 @@ def previous_seq_slots_device(): # For 3) 'dummy requests': pad dummy requests for CUDA graph or attention dp. # Already set to '0' during initialization. - num_extend_reqeust_wo_dummy = len(extend_requests) - len(extend_dummy_requests) self.previous_pos_id_offsets_cuda[ (num_extend_reqeust_wo_dummy - previous_batch_len) * runtime_tokens_per_gen_step : num_extend_reqeust_wo_dummy @@ -4189,8 +4220,7 @@ def previous_seq_slots_device(): # so we need to set the previous_pos_id_offsets and previous_kv_lens_offsets to zeros # to skip the value changes in _preprocess_inputs. Otherwise, there will be illegal memory access # when writing key/values to the KV cache. - self.previous_pos_id_offsets_cuda *= 0 - self.previous_kv_lens_offsets_cuda *= 0 + self._zero_previous_offsets_outside((0, 0), (0, 0)) # The host copy covers input_ids_cuda[:num_tokens] and the device copies # cover the previous-batch rows; padding-request rows are covered by @@ -4242,10 +4272,7 @@ def previous_seq_slots_device(): final_position_ids = self.position_ids_cuda[:total_num_tokens].unsqueeze(0) if enable_spec_decode: - self.gather_ids_cuda[: len(gather_ids)].copy_( - torch.tensor(gather_ids, dtype=torch.int, pin_memory=prefer_pinned()), - non_blocking=True, - ) + copy_if_changed(self.gather_ids_cuda[: len(gather_ids)], gather_ids) if self.mapping.has_cp_helix(): # A non-None owned-count list is what arms @@ -4487,9 +4514,6 @@ def previous_seq_slots_device(): # num_generations / num_tokens / seq_lens are set above, before the # attention-DP allgather that must agree with prepare(). spec_metadata.host_position_ids = host_position_ids - spec_metadata.num_accepted_draft_tokens = self.num_accepted_draft_tokens_cuda[ - : len(num_accepted_draft_tokens) - ] if context_prompt_lookahead is not None: spec_metadata.populate_context_prompt_lookahead(context_prompt_lookahead) # No-op for non 1-model diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder_decoder.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder_decoder.py index 8178deb67fea..60e3177295fb 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder_decoder.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder_decoder.py @@ -17,6 +17,7 @@ from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.attention.backends.vanilla import VanillaAttentionMetadata from tensorrt_llm._torch.distributed import Distributed +from tensorrt_llm._torch.host_staging import forget_copy from tensorrt_llm._torch.memory_buffer_utils import with_shared_pool from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import MoeLoadBalancer @@ -1350,6 +1351,8 @@ def _prepare_encoder_decoder_inputs_fast( buffers["previous_batch_indices"][:num_previous_batch_requests], non_blocking=True, ) + # The decoder staging writes this buffer through copy_if_changed(). + forget_copy(previous_slots) self._encoder_decoder_staged_request_ids = staged_request_ids generation_begin = num_context_tokens generation_end = generation_begin + num_previous_batch_requests diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index fc29a0633f41..bfc7ccba0a7f 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -16,6 +16,7 @@ from ..attention.backends import AttentionMetadata from ..attention.backends.flashinfer import FlashInferAttentionMetadata from ..attention.backends.sparse.params import MTPIndexShareMetadata +from ..host_staging import copy_if_changed from ..model_config import ModelConfig from ..pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.llm_request import LlmRequest @@ -348,12 +349,7 @@ def prepare(self): assert self.request_ids is not None # update batch indices num_seqs = len(self.request_ids) - batch_indices = torch.arange(num_seqs, - dtype=torch.int, - device='cpu', - pin_memory=prefer_pinned()) - self.batch_indices_cuda[:num_seqs].copy_(batch_indices, - non_blocking=True) + copy_if_changed(self.batch_indices_cuda[:num_seqs], range(num_seqs)) # Snapshot the real (unpadded) scheduled token count before the decrement # below. maybe_capture_hidden_states runs during the target forward and # must bound the capture by this count; self.num_tokens is about to be diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index 5fb5b185dc0b..02688203d157 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -28,6 +28,7 @@ import torch from ...sampling_params import SamplingParams +from ..host_staging import copy_if_changed from ..pyexecutor.llm_request import LlmRequest, LlmRequestState, get_draft_token_length from ..pyexecutor.resource_manager import BaseResourceManager from ..pyexecutor.sampler import ( @@ -166,6 +167,8 @@ class SpecSampler(Sampler[SampleStateSpec], AsyncWorkerMixin): """ SampleState = SampleStateSpec + # Store rows of the sampled requests in batch order; allocated on first use. + _sampling_slots: Optional[torch.Tensor] = None def is_generation_model(self) -> bool: return True @@ -442,8 +445,10 @@ def sample_async( for r in sampling_requests ] - slots = torch.as_tensor([r.py_seq_slot for r in sampling_requests], dtype=torch.long) - slots = slots.to(device="cuda", non_blocking=True) + if self._sampling_slots is None: + self._sampling_slots = torch.empty_like(self.store.new_tokens_lens, dtype=torch.long) + slots = self._sampling_slots[:num_sampling_requests] + copy_if_changed(slots, [r.py_seq_slot for r in sampling_requests]) o_new_tokens = outputs["new_tokens"][num_skip : num_skip + num_sampling_requests] o_new_tokens_lens = outputs["new_tokens_lens"][num_skip : num_skip + num_sampling_requests] From 8a8378b12d6a69e845cedeab1472558450237338 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Fri, 16 Oct 2026 10:00:29 -0700 Subject: [PATCH 52/53] [None][feat] Bound the speculative KV page-table copy and skip unchanged DSA page-table uploads KV cache manager V2 ignored max_blocks and always copied every block column of each row of its pinned page table (max_seq_len worth, about 16k columns at a 1M context) with copyBatchBlockOffsetsToDeviceKernel. The kernel and its binding now take an optional column bound (max_num_blocks, rounded up to whole 16-byte accesses; negative or absent copies whole rows), and KVCacheManagerV2.copy_batch_block_offsets passes max_blocks through as KVCacheManager already does. Helix rows are still copied whole. TrtllmAttentionMetadata.prepare() capped the copy only without speculative decoding, because the overlap correction and the draft sub-steps advance kv_lens_cuda on the device past the host lengths. With a linear draft each of these advances is below one verification window of 1 + max_total_draft_tokens tokens: the correction adds the previous step's accepted count minus this step's window, and the draft steps add at most draft_len - 1 after resetting to the accepted length. The cap under linear speculation is therefore the host maximum KV length (which already includes the extra draft KV tokens) plus two windows. Tree drafts and helix keep the full width. No kernel reads a block column past its device KV length, so the columns left unwritten are never read. The DSA INDEX_KEY offsets and block table were uploaded every step: - The host INDEX_KEY rows are now compared with the freshly computed page indices and rewritten only when they differ; each rewrite gives the host table a new version object. The upload of the batch rows goes through copy_if_changed() with that version as the token, so it is skipped while the device rows hold that version of the host rows. - The block table goes through copy_if_changed() with its host payload. Both device buffers are written only by these uploads, so a skipped upload leaves exactly the contents the copy would have written. The INDEX_KEY rows only change when a request gets a new block, and the block table only when its used width or entries change. Signed-off-by: Dongfeng Yu --- .../batch_manager/kvCacheManagerV2Utils.cu | 28 +++++++++++++------ .../batch_manager/kvCacheManagerV2Utils.h | 4 ++- .../batch_manager/kvCacheManagerV2Utils.cpp | 7 +++-- .../attention/backends/sparse/dsa/metadata.py | 26 ++++++++++++----- .../_torch/attention/backends/trtllm.py | 26 +++++++++++------ .../kv_cache/kv_cache_manager_v2.py | 11 ++++++-- 6 files changed, 73 insertions(+), 29 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cu b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cu index af13ae72a61c..6545a3a4da23 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cu +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cu @@ -184,19 +184,20 @@ CUresult copyDeviceToDevice(std::vector const& tasks, ssize_t numBytes, return launchBatchedCopy(false, tasks, numBytes, stream); } -// dst_tensor[:, :num_seqs, 0] = src_tensor[:, copy_idx] -// dst_tensor[:, :num_seqs, 1] = dst_tensor[:, :num_seqs, 0] + 1 +// dst_tensor[:, :num_seqs, 0, :numBlocksToCopy] = src_tensor[:, copy_idx, 0, :numBlocksToCopy] +// dst_tensor[:, :num_seqs, 1, :numBlocksToCopy] = dst_tensor[:, :num_seqs, 0, :numBlocksToCopy] + 1 +// Rows are numBlocksPerSeq long in both tensors; the columns past numBlocksToCopy are not written. __global__ void copyBatchBlockOffsetsToDeviceKernel(SizeType32 const* __restrict__ srcPtr, SizeType32* __restrict__ dstPtr, SizeType32 const srcMaxNumSequences, SizeType32 const dstMaxNumSequences, - SizeType32 numBlocksPerSeq, SizeType32 const* __restrict__ copyIndex, SizeType32 const* __restrict__ indexScales, - SizeType32 const* __restrict__ kvOffset) + SizeType32 numBlocksPerSeq, SizeType32 numBlocksToCopy, SizeType32 const* __restrict__ copyIndex, + SizeType32 const* __restrict__ indexScales, SizeType32 const* __restrict__ kvOffset) { constexpr uint32_t kvFactor = 2; constexpr auto elemPerAccess = sizeof(PackedInt) / sizeof(SizeType32); __shared__ PackedInt data[copyBlocknbBufs][copyBlockCtaSize]; - auto const iterPerSeq = divUp(numBlocksPerSeq * sizeof(SizeType32), sizeof(PackedInt) * copyBlockCtaSize); + auto const iterPerSeq = divUp(numBlocksToCopy * sizeof(SizeType32), sizeof(PackedInt) * copyBlockCtaSize); auto const tid = threadIdx.x; auto const poolIdx = blockIdx.x; auto const seqIdx = blockIdx.y; @@ -205,7 +206,7 @@ __global__ void copyBatchBlockOffsetsToDeviceKernel(SizeType32 const* __restrict uint32_t const dstIdxKBeg = tid * elemPerAccess + (poolIdx * dstMaxNumSequences + seqIdx) * seqDimStride; uint32_t const dstIdxVBeg = dstIdxKBeg + numBlocksPerSeq; - uint32_t const srcIdxEnd = (poolIdx * srcMaxNumSequences + copyIndex[seqIdx]) * seqDimStride + numBlocksPerSeq; + uint32_t const srcIdxEnd = (poolIdx * srcMaxNumSequences + copyIndex[seqIdx]) * seqDimStride + numBlocksToCopy; for (uint32_t i = 0; i < iterPerSeq + copyBlocknbBufs; i++) { @@ -250,7 +251,7 @@ __global__ void copyBatchBlockOffsetsToDeviceKernel(SizeType32 const* __restrict // Host-side launcher void copyBatchBlockOffsetsToDevice(ITensor const& input, ITensor& output, ITensor const& copyIndex, - ITensor const& indexScales, ITensor const& kvOffset, CUstream stream) noexcept + ITensor const& indexScales, ITensor const& kvOffset, SizeType32 maxNumBlocks, CUstream stream) noexcept { using namespace tensorrt_llm::runtime; @@ -282,11 +283,22 @@ void copyBatchBlockOffsetsToDevice(ITensor const& input, ITensor& output, ITenso "Not implemented case: numBlocksPerSeq * sizeof(SizeType32) = %zu must be a multiple of %zu.", static_cast(numBlocksPerSeq * sizeof(SizeType32)), static_cast(sizeof(PackedInt))); + // A negative maxNumBlocks copies whole rows. Otherwise the copy covers the first maxNumBlocks + // columns, rounded up to whole PackedInt accesses, which stay inside the row. + SizeType32 numBlocksToCopy = numBlocksPerSeq; + if (maxNumBlocks >= 0) + { + constexpr SizeType32 blocksPerAccess = sizeof(PackedInt) / sizeof(SizeType32); + SizeType32 const roundedBlocks + = (std::max(maxNumBlocks, SizeType32{1}) + blocksPerAccess - 1) / blocksPerAccess * blocksPerAccess; + numBlocksToCopy = std::min(numBlocksPerSeq, roundedBlocks); + } + dim3 gridDim(numPools, numSeqs, 1); dim3 blockDim(copyBlockCtaSize); copyBatchBlockOffsetsToDeviceKernel<<>>(srcPtr, dstPtr, srcMaxNumSequences, - dstMaxNumSequences, numBlocksPerSeq, copyIndexPtr, indexScalesPtr, kvOffsetPtr); + dstMaxNumSequences, numBlocksPerSeq, numBlocksToCopy, copyIndexPtr, indexScalesPtr, kvOffsetPtr); } } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h index f15e96daf10b..e8f4308f120f 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h @@ -108,7 +108,9 @@ CUresult copyHostToDevice(std::vector> const& tasks CUresult copyDeviceToHost(std::vector> const& tasks, ssize_t numBytes, CUstream stream); CUresult copyDeviceToDevice(std::vector> const& tasks, ssize_t numBytes, CUstream stream); +//! Copies the block offsets of the batch rows. A non-negative maxNumBlocks bounds the copied columns +//! (rounded up to a multiple of four); a negative one copies whole rows. void copyBatchBlockOffsetsToDevice(ITensor const& input, ITensor& output, ITensor const& copyIndex, - ITensor const& indexScales, ITensor const& kvOffset, CUstream stream) noexcept; + ITensor const& indexScales, ITensor const& kvOffset, SizeType32 maxNumBlocks, CUstream stream) noexcept; } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.cpp index d9bbee43305a..4a725a133284 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.cpp @@ -138,7 +138,7 @@ void KVCacheManagerV2UtilsBindings::initBindings(nb::module_& module) module.def( "copy_batch_block_offsets_to_device", [](at::Tensor input, at::Tensor output, at::Tensor copyIndex, at::Tensor indexScales, at::Tensor kvOffset, - uintptr_t stream) + uintptr_t stream, SizeType32 maxNumBlocks) { auto _input = from_torch(input); auto _output = from_torch(output); @@ -151,10 +151,11 @@ void KVCacheManagerV2UtilsBindings::initBindings(nb::module_& module) TLLM_CHECK_WITH_INFO(_indexScales.has_value(), "Invalid index scales tensor."); TLLM_CHECK_WITH_INFO(_kvOffset.has_value(), "Invalid kv offset tensor."); copyBatchBlockOffsetsToDevice(*(_input.value()), *(_output.value()), *(_copyIndex.value()), - *(_indexScales.value()), *(_kvOffset.value()), reinterpret_cast(stream)); + *(_indexScales.value()), *(_kvOffset.value()), maxNumBlocks, reinterpret_cast(stream)); }, nb::arg("input"), nb::arg("output"), nb::arg("copy_index"), nb::arg("index_scales"), nb::arg("kv_offset"), - nb::arg("stream"), nb::call_guard(), "Copy batch block indices to device"); + nb::arg("stream"), nb::arg("max_num_blocks") = -1, nb::call_guard(), + "Copy batch block indices to device; a non-negative max_num_blocks bounds the copied columns"); } } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py index ae4188c7b20b..736d74e417af 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py @@ -11,11 +11,13 @@ from typing import TYPE_CHECKING, ClassVar, FrozenSet, List, Optional, Tuple import torch +from torch.utils.weak import WeakIdKeyDictionary import tensorrt_llm import tensorrt_llm.bindings from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE +from tensorrt_llm._torch.host_staging import copy_if_changed from tensorrt_llm._torch.utils import maybe_compile from tensorrt_llm._utils import get_sm_version, prefer_pinned from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata @@ -282,6 +284,9 @@ def __init__(self, *args, **kwargs): self._radix_rows_per_sequence = 1 # Set only inside defer_cuda_graph_decode_outputs(). self._defer_decode_outputs = False + # Per INDEX_KEY host table, an object replaced whenever the table's + # contents change; device uploads of the table carry it as a token. + self._index_key_host_versions = WeakIdKeyDictionary() super().__init__(*args, **kwargs) sparse_metadata_params = self.sparse_metadata_params if not isinstance(sparse_metadata_params, DSAMetadataParams): @@ -1499,9 +1504,18 @@ def _update_indexer_k_cache_block_offsets(self) -> torch.Tensor: ) page_indices = pool_indices * cache_manager.indexer_k_cache_page_scale num_blocks = page_indices.shape[1] - self.host_indexer_k_cache_block_offsets[: self.num_seqs, :num_blocks].copy_(page_indices) - self.indexer_k_cache_block_offsets[: self.num_seqs].copy_( - self.host_indexer_k_cache_block_offsets[: self.num_seqs], non_blocking=True + host_offsets = self.host_indexer_k_cache_block_offsets + host_rows = host_offsets[: self.num_seqs, :num_blocks] + version = self._index_key_host_versions.get(host_offsets) + if version is None or not torch.equal(host_rows, page_indices): + host_rows.copy_(page_indices) + version = self._index_key_host_versions[host_offsets] = object() + # Rows the device already took from this version of the host table + # are not uploaded again. + copy_if_changed( + self.indexer_k_cache_block_offsets[: self.num_seqs], + host_offsets[: self.num_seqs], + token=version, ) # The batch rows hold non-negative pool slots. Sanitize the # graph-padding rows after them, which may be stale after cache @@ -2072,10 +2086,8 @@ def prepare_for_indexer_k_cache(self): host_block_table.masked_fill_( pad_cols.unsqueeze(0) >= num_blocks_per_seq[: self.num_seqs].unsqueeze(1), -1 ) - # Copy to GPU - self.block_table[: self.num_seqs, :max_blocks_used].copy_( - host_block_table, non_blocking=True - ) + # Copy to GPU unless the device rows already hold this table. + copy_if_changed(self.block_table[: self.num_seqs, :max_blocks_used], host_block_table) def prepare_for_skip_indexer(self, kv_lens: torch.Tensor, dense_topk_indices: bool = True): """Decide which requests skip the indexer. diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index d29def811bc5..1cbe4d393a8c 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -1077,20 +1077,30 @@ def prepare(self) -> None: # On the non-speculative path the host kv_lens snapshot bounds # every block-table access, so the staged/H2D width can be capped # at the batch's maximum instead of max_seq_len's worth of - # columns. Speculative decoding must stage the full width: - # draft/tree sub-steps and the overlap scheduler advance - # kv_lens_cuda on device past the host snapshot, and their - # kernels dereference block columns a host-derived cap would - # leave unstaged (uninitialized in this buffer). + # columns. Under speculative decoding the overlap scheduler and + # the draft sub-steps advance kv_lens_cuda on device past the host + # snapshot, and their kernels dereference block columns a + # host-derived cap would leave unstaged (uninitialized in this + # buffer). With a linear draft each of the two advances is less + # than one verification window of 1 + max_total_draft_tokens + # tokens, so the cap keeps two windows of room. Tree drafts and + # helix stage the full width. spec_active = (self.draft_kv_cache_manager is not None or self.is_spec_decoding_enabled or bool(self.kv_cache_params.num_extra_kv_tokens) or (self.runtime_features is not None and self.runtime_features.has_speculative_draft_tokens)) max_blocks = None - if not spec_active and self.kv_cache_manager.tokens_per_block: - max_blocks = ceil_div(max_kv_len, - self.kv_cache_manager.tokens_per_block) + tokens_per_block = self.kv_cache_manager.tokens_per_block + if tokens_per_block and not spec_active: + max_blocks = ceil_div(max_kv_len, tokens_per_block) + elif (tokens_per_block and self.max_total_draft_tokens is not None + and not self.is_spec_dec_tree + and not self.is_spec_dec_dynamic_tree + and not self.enable_helix): + max_blocks = ceil_div( + max_kv_len + 2 * (1 + self.max_total_draft_tokens), + tokens_per_block) self.kv_cache_manager.copy_batch_block_offsets( self.kv_cache_block_offsets, self.request_ids, diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 9aa7d619a3e8..f1aad6acd9da 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -6536,8 +6536,6 @@ def copy_batch_block_offsets( "Offloaded sparse history requires Batch metadata and sparse fetch; " "dense attention offsets cannot address host slots" ) - # max_blocks is accepted for signature parity with KVCacheManager; the - # device-side copy op here already scales with allocated blocks only. assert beam_width <= self.max_copy_beam_width # Cross KV is request-scoped: generation beams share the encoder K/V @@ -6556,6 +6554,14 @@ def copy_batch_block_offsets( ) return + # As in KVCacheManager, max_blocks bounds the copied block columns of + # each row; the columns past it keep their previous contents. Helix + # rows are copied whole. + bound = ( + {"max_num_blocks": max_blocks} + if max_blocks is not None and not self._has_cp_helix + else {} + ) copy_batch_block_offsets_to_device( self.host_kv_cache_block_offsets, dst_tensor, @@ -6563,6 +6569,7 @@ def copy_batch_block_offsets( self.index_scales, self.kv_offset, self._stream.cuda_stream, + **bound, ) @staticmethod From 3ee542d591ec25c053046d786817cba1331ebae0 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Fri, 16 Oct 2026 10:01:45 -0700 Subject: [PATCH 53/53] [None][feat] Upload the per-request Philox state only for rejection sampling SpecMetadata._populate_request_rng_state() uploaded the per-request seeds and offsets (request_seeds, request_offsets) on every step. Their only reader is _rng_state_per_request(), which runs only on the rejection-sampling paths: the rejection acceptance kernels and the advanced draft samplers, which are taken only when use_rejection_sampling is set and the batch is not all greedy. use_rejection_sampling is fixed when the metadata is created, so without it the two H2D copies are skipped. The host offset windows are still taken every step, so the per-row seeds and offsets the non-greedy target sampling reads are unchanged. Signed-off-by: Dongfeng Yu --- tensorrt_llm/_torch/speculative/interface.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 8237c09ee8b7..ba1d1b8b8ad8 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -836,7 +836,8 @@ def _populate_request_rng_state( Both layouts are produced: ``request_*`` with one entry per request, and ``seeds`` / ``offsets`` expanded to one entry per logits row (the temperatures / top_ks / top_ps / min_ps layout), because the sampling calls - take one or the other. + take one or the other. Only rejection sampling and its draft sampling + read ``request_*``, so without it those buffers are not uploaded. A request that specified no seed gets ``DEFAULT_SAMPLING_SEED`` and an offset window from the counter shared by all unseeded requests. Each of @@ -899,8 +900,9 @@ def _upload(dst: torch.Tensor, values: list[int]) -> None: pin_memory=prefer_pinned()), non_blocking=True) - _upload(self.request_seeds, request_seeds) - _upload(self.request_offsets, request_offsets) + if self.use_rejection_sampling: + _upload(self.request_seeds, request_seeds) + _upload(self.request_offsets, request_offsets) _upload(self.seeds, flat_seeds) _upload(self.offsets, flat_offsets)