diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 26995cec1b90..549a9516b47c 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -524,10 +524,15 @@ /tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency @dongfengy @Tabrizian @xwang233 /tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency @dongfengy @Tabrizian @xwang233 -# GLM-5.3-Flash (glm5_next) -/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next @zhaoyangwang-nvidia -/tests/unittest/_torch/modeling_v2/models/glm5_next @zhaoyangwang-nvidia -/tests/integration/defs/accuracy/modeling_v2/glm5_next/ @zhaoyangwang-nvidia +# GLM-5.3-Flash family (glm5_next): routing and shared code are shared by its target owners. +/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next @zhaoyangwang-nvidia @ruocheng-nv +/tests/unittest/_torch/modeling_v2/models/glm5_next @zhaoyangwang-nvidia @ruocheng-nv +/tests/integration/defs/accuracy/modeling_v2/glm5_next/ @zhaoyangwang-nvidia @ruocheng-nv + +# GLM-5.3-Flash block-FP8 sm_10x tp4 target: source, ops, native code, and tests. +/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4 @ruocheng-nv +/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4 @ruocheng-nv +/tests/integration/defs/accuracy/modeling_v2/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4 @ruocheng-nv # Uncomment only on release branches (e.g. release/0.19): # * @NVIDIA/trt-llm-release-branch-approval diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/__init__.py index b79138258c3a..bdaee68830e2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/__init__.py @@ -1,3 +1,3 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Glm5NextForConditionalGeneration targets.""" +"""Glm5NextForConditionalGeneration targets (GLM-5.3-Flash) and the routing between them.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/csrc/CMakeLists.txt b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/csrc/CMakeLists.txt new file mode 100644 index 000000000000..b471ee5a44ac --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/csrc/CMakeLists.txt @@ -0,0 +1,22 @@ +# 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. + +add_modeling_v2_family(glm5_next) + +add_modeling_v2_target( + glm5_next glm_5_3_flash_fp8__sm_10x__tp4 + ARCHS 100f + SOURCES *.cu *.cpp + REQUIRES MULTI_DEVICE) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/README.md new file mode 100644 index 000000000000..ff9a195a37fe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/README.md @@ -0,0 +1,91 @@ + + + +# GLM-5.3-Flash FP8 on 4x B200: deployment + +This target (`glm_5_3_flash_fp8__sm_10x__tp4`) serves the GLM-5.3-Flash FP8 checkpoint on one node of four +B200 (sm_100) GPUs with tensor parallelism 4 and MTP-3 speculative decoding. Routing does not select it on +sm_103 (B300 / GB300), where it has not run. Compared with a standard +GLM-5.3-Flash deployment, it needs `modeling_v2: require` in the server config and a larger `--max_num_tokens`; its +CUDA kernels are in the glm5_next family library the build compiles. + +## 1. Build + +Build or install TensorRT LLM as usual, for a CUDA architecture list that includes `100` and with multi-device +support (the default): the build then compiles the target's kernels into `libth_modeling_v2_glm5_next.so`. + +## 2. Hardware and checkpoint + +- One node, 4x B200, NVLink. +- The GLM-5.3-Flash FP8 checkpoint (128 x 128 block-FP8 experts and MLA projections). + +## 3. Server config + +`config.yaml`: + +```yaml +modeling_v2: require # select this ModelingV2 target; startup fails if the setup does not match it +speculative_config: + decoding_type: MTP + max_draft_len: 3 # one-model MTP-3 on the checkpoint's MTP layer +kv_cache_config: + dtype: fp8 + free_gpu_memory_fraction: 0.5 + enable_block_reuse: false +cuda_graph_config: + enable_padding: true + max_batch_size: 64 +enable_chunked_prefill: true +enable_attention_dp: false +disable_mm_encoder: true +moe_config: + backend: AUTO +``` + +`free_gpu_memory_fraction` can be raised for more KV-cache capacity. + +## 4. Launch + +```bash +trtllm-serve /path/to/GLM-5.3-Flash \ + --backend pytorch --tp_size 4 --ep_size 1 \ + --max_batch_size 64 --max_num_tokens 32768 --max_seq_len 131072 \ + --trust_remote_code --config config.yaml +``` + +What differs from a standard deployment: + +- `--max_num_tokens 32768`: long prompts prefill in larger chunks. The measurements below use it; the default + token budget has not been measured on this target. +- `--max_seq_len 131072` is the model's context length. Prompts of up to 13,312 tokens have been measured on this + target (the operations' certified cells stop there); longer ones are untested. +- Without the glm5_next family library (a build whose CUDA architectures include no `100`, or without multi-device + support) the target is not selected, and `modeling_v2: require` fails at startup. +- Keep the all-reduce strategy at its default (AUTO, or ONESHOT): the target runs the one-shot all-reduce, whose + workspace its fused kernels share, and routing declines other strategies. + +## 5. Check + +- The server log shows `modeling_v2: ModelingV2Glm53FlashFp8Sm10xTp4 built`. + +## 6. Measured + +One node of 4x B200, the config and launch above, against the built-in GLM-5.3-Flash model (`modeling_v2: off`) +with the same settings on the same node. Throughput: 13,312-token prompts, 33 output tokens (ignore_eos, +temperature 0), the median of two runs of 32 requests at C1 and 64 or 128 at C8 / C32 / C64. + +| | built-in | this target | +| --- | --- | --- | +| C1 TTFT p50 | 271 ms | 205 ms | +| C1 output tok/s/user p50 | 397 | 492 | +| C8 output tok/s | 126 | 175 | +| C32 output tok/s | 135 | 193 | +| C64 output tok/s | 135 | 198 | +| KV-cache memory per rank (`free_gpu_memory_fraction` 0.5) | 45.2 GiB | 44.0 GiB | +| GSM8K, accuracy harness, no speculation | 93.2% | 92.5% | +| GSM8K, accuracy harness, MTP-3 | 92.3% | 93.1% | +| MTP-3 acceptance length (accuracy harness) | 3.382 | 3.380 | +| GSM8K, 400 questions through the server, MTP-3 | 98.5% | 99.0% | + +The accuracy-harness rows are `test_accuracy.py` (5-shot, no chat template, the CI protocol); the last row asks +the server above zero-shot through its OpenAI chat endpoint (greedy, chat template), a different protocol. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/__init__.py new file mode 100644 index 000000000000..3bb60a1b1ed7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/__init__.py @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GLM-5.3-Flash block-FP8 on sm_100: tensor-parallel attention and experts over 4 ranks.""" + +from tensorrt_llm._torch._experimental.modeling_v2._native import load_target_ops + +native_ops = load_target_ops(__name__) # torch.ops.mv2_glm5_next_glm_5_3_flash_fp8_sm_10x_tp4 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/checkpoint.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/checkpoint.py new file mode 100644 index 000000000000..934f59c6cfce --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/checkpoint.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Reads this rank's share of checkpoint tensors and records what was read. + +The engine hands ``load_weights`` a mapping of checkpoint keys to tensors or lazy +safetensors slices. ``Checkpoint.read`` cuts a lazy entry before it is materialized, so a +rank reads only its rows or columns of a split weight, and records every key it touches: +after loading, every key must have been read except the ones the target predicts it does +not load (the vision tower, and the MTP layer without speculative decoding). +""" + +from __future__ import annotations + +from typing import Any, Iterable, Mapping, Optional + +import torch + +#: An index into a checkpoint tensor: one slice per leading axis. +Cut = tuple[slice, ...] + +#: The side of the checkpoint's square FP8 scale blocks (``weight_scale_inv``). +BLOCK = 128 + + +class Checkpoint: + def __init__(self, weights: Mapping[str, Any], device: torch.device) -> None: + self.weights = weights + self.device = device + self.read_keys: set[str] = set() + + def read( + self, key: str, cut: Optional[Cut] = None, dtype: Optional[torch.dtype] = None + ) -> torch.Tensor: + """``key``'s tensor (``cut`` of it, read from the lazy entry), on the device. + + ``dtype`` converts after the read; ``None`` keeps the checkpoint's dtype (FP8 + payloads and FP32 block scales stay as stored). + """ + if key not in self.weights: + raise KeyError(f"checkpoint has no tensor {key!r}") + entry = self.weights[key] + if cut is not None: + tensor = entry[cut] + else: + tensor = entry if torch.is_tensor(entry) else entry[...] + self.read_keys.add(key) + tensor = tensor.to(device=self.device) + if dtype is not None and tensor.dtype != dtype: + tensor = tensor.to(dtype) + return tensor.contiguous() + + def check_coverage(self, skipped: Iterable[str]) -> None: + """Every key was read or is predicted to be skipped; nothing predicted was read.""" + skipped = set(skipped) + keys = set(self.weights.keys()) + unread = sorted(keys - self.read_keys - skipped) + if unread: + raise ValueError(f"{len(unread)} checkpoint tensors were not loaded, e.g. {unread[:5]}") + both = sorted(self.read_keys & skipped) + if both: + raise ValueError(f"tensors predicted to be skipped were loaded: {both[:5]}") + + +def rows(start: int, end: int) -> Cut: + """The cut of rows ``[start, end)``.""" + return (slice(start, end),) + + +def cols(start: int, end: int) -> Cut: + """The cut of columns ``[start, end)`` of a 2-D tensor.""" + return (slice(None), slice(start, end)) + + +def share(total: int, tp_size: int, tp_rank: int) -> tuple[int, int]: + """This rank's contiguous ``[start, end)`` of ``total`` split evenly over ``tp_size``.""" + if total % tp_size: + raise ValueError(f"{total} does not split evenly over {tp_size} ranks") + size = total // tp_size + return tp_rank * size, (tp_rank + 1) * size + + +def dequantize_block_fp8( + weight: torch.Tensor, scale: torch.Tensor, block: int = BLOCK +) -> torch.Tensor: + """BF16 ``weight`` from a 128 x 128 block-FP8 weight and its FP32 block scales.""" + n, k = weight.shape + assert scale.shape == (-(-n // block), -(-k // block)), ( + f"block scales {tuple(scale.shape)} for a {n} x {k} weight" + ) + full = scale.float().repeat_interleave(block, 0)[:n].repeat_interleave(block, 1)[:, :k] + return (weight.float() * full).to(torch.bfloat16) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/csrc/decodeMoe.cu b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/csrc/decodeMoe.cu new file mode 100644 index 000000000000..f0d9a735014c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/csrc/decodeMoe.cu @@ -0,0 +1,1072 @@ +/* + * 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. + */ +// Decode MoE (M <= 4 rows): one persistent warp-specialized kernel per call, one CTA per SM. +// +// Shapes (GLM-5.3-Flash at moe_tp4): H = 4096, I = 512, E = 288, top-8 + the shared expert. +// w31 [E, 2I, H] e4m3 = [up; gate] rows, s31 [E, 2I/128, H/128]; w2 [E, H, I] e4m3, s2 [E, H/128, I/128]. +// In the code, M is numRows, G (the grid) numCtas and U (the expert slots) numSlots. +// +// Every CTA of the grid must be resident at once: FC2 CTAs spin on counters that other CTAs' FC1 epilogues fill, +// and the shared memory admits one CTA per SM. decodeMoe() checks that the device can hold the grid and returns +// cudaErrorCooperativeLaunchTooLarge otherwise; SMs held by other work at launch time are not visible to it. +// +// The op is a stream of U x 6.29 MB of FP8 weights (U = distinct experts, the shared one included). Every CTA +// computes the routing itself and derives the same static schedule from it (no work counter, no inter-CTA +// hand-off of the routing): +// FC1: a tile = 4 up + 4 gate rows (all of H, 32 KB) of one expert, global tile order (expert slot, tile), the +// shared expert first. Tiles [0, tl) go round robin over the G CTAs (tile j -> CTA j % G); the tail +// [tl, U x 128) goes to the G - 128 CTAs without FC2 rows, about the bytes of the FC2 work they skip, less +// a margin so that the last experts' activations are out before FC2 needs them. +// FC2: CTA c < 128 owns output rows [32 c, 32 c + 32) and walks all U experts (two per stage), summing them in +// a fixed order (deterministic, no atomics on the output). +// Routing: warp t < M takes token row t, top-8 by 8 rounds of a warp max (the winner lane drops its key); warp 0 +// then numbers the distinct experts in first-occurrence order with match.any (slot 0 = the shared expert). Before +// griddepcontrol.wait (only weights are read there) CTAs c < 128 issue the shared expert's FC1 tile c and +// prefetch 16 KB of its w2 into L2; the first stages after row 0's top-8 are issued before the other rows finish. +// Warp roles: 8 MMA consumer warps, one producer warp (the CTA's whole stage list, FC1 then FC2 weights, through +// a TMA bulk ring: FC2 weights are fetched while FC1 still computes), one act-loader warp, kNEpi FC1-epilogue +// warps. An FC1 tile's epilogue sums the consumer warps' K slices, applies the clamped SwiGLU and the routing +// weight, publishes the M activation rows (f16 hi + lo) to global scratch and bumps the expert slot's counter +// (release); the act loader waits for counter == 128 (acquire) and copies the activation into a small act ring. +// The counters are double-buffered by call parity (per-CTA call epochs): a call zeroes the set the next call uses. +// +// Arithmetic: e4m3 -> f16 (exact, cvt.rn.f16x2.e4m3x2), BF16 x -> f16, mma.sync m16n8k16 f16 -> f32, every +// 128-wide K block's partial scaled by its block scale in f32. +// K permutation: in a 64-wide K block, lane (g = lane / 4, t4 = lane % 4) reads bytes [16 t4, 16 t4 + 16) of +// its weight rows; MMA chunk j takes bytes 16 t4 + 4 j + {0, 1} as A columns {2 t4, 2 t4 + 1} and + {2, 3} as +// {2 t4 + 8, 2 t4 + 9}; the B fragment (column g = token) reads the activation at the same K positions. +#include "glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/csrc/decodeMoe.h" + +#include +#include +#include + +#include + +TRTLLM_MV2_NAMESPACE_BEGIN + +namespace kernels +{ +namespace +{ + +constexpr int kH = kDecodeMoeHidden, kI = kDecodeMoeInter, kE = kDecodeMoeExperts, kTopK = 8; +constexpr int kMaxM = kDecodeMoeMaxRows; +constexpr int kMaxU = 1 + kTopK * kMaxM; // 33 +constexpr int kTilesPerExpert = kI / 4; // 128 FC1 tiles of 4 up + 4 gate rows +constexpr int kCWarps = 8, kConsumers = kCWarps * 32; +constexpr int kProducerWarp = kCWarps, kActWarp = kCWarps + 1, kEpiWarp = kCWarps + 2; +constexpr int kNEpi = 6; // FC1 epilogue warps (round robin over tiles: their release fences overlap) +constexpr int kThreads = kConsumers + 64 + 32 * kNEpi; +constexpr int kNS = 5; // weight ring slots +constexpr int kEarly = kNS; // FC1 stages issued right after row 0's top-8 (expert slots <= 5) +constexpr int kEarlyPicks = 5; // row 0's picks those stages need (slots 1..5) +constexpr int kFc2Ctas = 128; // CTAs with FC2 rows (32 each) +constexpr int kMinFc1OnlyCtas = 4; // the fewest CTAs without FC2 rows the tail split runs with +// The largest grid: B200's 148 SMs, the grid the schedule (kTailMargin, kEarly) was tuned and certified on. The +// early stages need it to stay small: they are the CTA's tiles c + s G, s < kEarly, issued once row 0's first +// kEarlyPicks picks are known, so those tiles must lie in expert slots <= kEarlyPicks. +constexpr int kMaxCtas = 148; +constexpr int kTailMargin = 10; // FC1 tiles kept off the FC1-only CTAs' tail +constexpr int kSlotW = 32768; // weight bytes per slot +constexpr int kSlotBytes = kSlotW + 256; // + scales +constexpr int kNA = 7; // activation ring slots (the rest of the shared memory) +constexpr int kActPitch = 520, kActPlane = kMaxM * kActPitch; +constexpr int kActBytes = 2 * kActPlane * 2; // padded [hi, lo][kMaxM][kActPitch] f16 +constexpr int kRed2Stride = 5; +constexpr int kFc2Pitch = 16384; // one expert's <= 32 rows x 512 in a slot +constexpr int kMaxRows2 = 32; +constexpr int kCntPad = 64; // counters 256 B apart +constexpr int kRed1Floats = kCWarps * 16 * 2; // FC1 hand-off: [warp][g * 2 + t4 (< 2)][2] +constexpr int kSmem = kNS * kSlotBytes + kNA * kActBytes + kNEpi * kRed1Floats * 4; +// FC1: a consumer warp's K slice (of all of H), its 64-wide blocks +constexpr int kKW1 = kH / kCWarps, kKB1 = kKW1 / 64; +static_assert(kMaxM * kH * 2 <= kNA * kActBytes, "x is staged in the act ring"); +// FC2: warp = (K part kp, m-tile mt), K part width, its 64-wide blocks +static_assert(kKB1 == 8 && kCWarps == 8, "warp split (FC1: 512-wide K slices; FC2: 64-wide K parts)"); +static_assert(kCWarps * 32 * kRed2Stride <= kNEpi * kRed1Floats, "FC2 K-part reduction fits the FC1 hand-off buffers"); +static_assert(kEarly * kMaxCtas <= (kEarlyPicks + 1) * kTilesPerExpert, + "the early stages' tiles lie in the expert slots row 0's first kEarlyPicks picks fill"); +static_assert(kFc2Ctas + kMinFc1OnlyCtas <= kMaxCtas, "the smallest grid fits under the largest"); + +//! The kernel's arguments: DecodeMoeArgs' operands plus the device-global scratch. +struct MoeParams +{ + __nv_bfloat16 const* x; //!< [M, H] + float const* logits; //!< [M, E] + float const* bias; //!< [E] + uint8_t const* w31; //!< [E, 2I, H] e4m3, [up; gate] + float const* s31; //!< [E, 2I/128, H/128] + uint8_t const* sw31; //!< [2I, H] e4m3, the shared expert + float const* ss31; //!< [2I/128, H/128] + uint8_t const* w2; //!< [E, H, I] e4m3 + float const* s2; //!< [E, H/128, I/128] + uint8_t const* sw2; //!< [H, I] e4m3, the shared expert + float const* ss2; //!< [H/128, I/128] + __half* act; //!< [kMaxU][2][kMaxM][I] + //! [2][kMaxU] finished FC1 tiles per expert slot: set (call parity) of this call, the other is reset + int* cnt; + unsigned* epoch; //!< [CTA] calls so far (every CTA runs every call: they agree on the parity) + __nv_bfloat16* out; //!< [M, H] + float scaling; //!< the routed scaling factor + float limit; //!< the SwiGLU clamp + int numRows; //!< M +}; + +__device__ __forceinline__ uint32_t e4m3x2ToF16x2(uint32_t v16) +{ + uint32_t r; + uint16_t const h = static_cast(v16); + asm("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(r) : "h"(h)); + return r; +} + +__device__ __forceinline__ void mma16816( + 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)); +} + +__device__ __forceinline__ uint32_t word(uint4 const& v, int j) +{ + return j == 0 ? v.x : j == 1 ? v.y : j == 2 ? v.z : v.w; +} + +__device__ __forceinline__ uint32_t smemU32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ void mbarInit(uint64_t* bar, int count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" ::"r"(smemU32(bar)), "r"(count)); +} + +__device__ __forceinline__ void mbarInitFence() +{ + asm volatile("fence.mbarrier_init.release.cluster;" ::: "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 mbarArrive(uint64_t* bar) +{ + asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" ::"r"(smemU32(bar)) : "memory"); +} + +__device__ __forceinline__ void mbarWait(uint64_t* bar, uint32_t parity) +{ + asm volatile( + "{\n .reg .pred p;\n WAIT_%=:\n mbarrier.try_wait.parity.shared::cta.b64 p, [%0], %1;\n @!p bra WAIT_%=;\n}" :: + "r"(smemU32(bar)), + "r"(parity) + : "memory"); +} + +__device__ __forceinline__ uint64_t evictFirstPolicy() +{ + uint64_t pol; + asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;" : "=l"(pol)); + return pol; +} + +// TMA bulk copy global -> shared, completing `bytes` on bar +// weights: streamed once, evict-first in L2 (keeps the prefetched / published data there) +__device__ __forceinline__ void bulkG2s(void* dst, void const* src, uint32_t bytes, uint64_t* bar, uint64_t pol) +{ + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], %4;" :: + "r"(smemU32(dst)), + "l"(src), "r"(bytes), "r"(smemU32(bar)), "l"(pol) + : "memory"); +} + +__device__ __forceinline__ void bulkG2s(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__ void l2Prefetch(void const* src, uint32_t bytes) +{ + uint64_t pol; + asm volatile("createpolicy.fractional.L2::evict_last.b64 %0, 1.0;" : "=l"(pol)); + asm volatile("cp.async.bulk.prefetch.L2.global.L2::cache_hint [%0], %1, %2;" ::"l"(src), "r"(bytes), "l"(pol) + : "memory"); +} + +__device__ __forceinline__ int ldAcquire(int const* p) +{ + int v; + asm volatile("ld.acquire.gpu.global.b32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ void fenceProxyAsyncGlobal() +{ + asm volatile("fence.proxy.async.global;" ::: "memory"); +} + +__device__ __forceinline__ void redReleaseAdd(int* p, int v) +{ + asm volatile("red.release.gpu.global.add.s32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} + +__device__ __forceinline__ void consumerSync() +{ + asm volatile("bar.sync 1, %0;" ::"n"(kConsumers) : "memory"); +} + +// The largest key not in mask (lowest index on ties; 0 if none): a depth-4 tree. +__device__ __forceinline__ void bestOf9( + uint32_t const (&key)[9], float const (&sc)[9], unsigned mask, uint32_t& bk, int& bi, float& bs) +{ + uint32_t v[9]; + int ix[9]; + float sv[9]; +#pragma unroll + for (int i = 0; i < 9; i++) + { + v[i] = ((mask >> i) & 1u) ? 0u : key[i]; + ix[i] = i; + sv[i] = sc[i]; + } +#pragma unroll + for (int st = 1; st < 16; st <<= 1) + { +#pragma unroll + for (int i = 0; i + st < 9; i += 2 * st) + { + bool const hi = v[i + st] > v[i]; // ties keep the lower index i + v[i] = hi ? v[i + st] : v[i]; + ix[i] = hi ? ix[i + st] : ix[i]; + sv[i] = hi ? sv[i + st] : sv[i]; + } + } + bk = v[0]; + bi = ix[0]; + bs = sv[0]; +} + +__device__ __forceinline__ uint32_t f2ord(float f) +{ + uint32_t const b = __float_as_uint(f); + return (b & 0x80000000u) ? ~b : (b | 0x80000000u); +} + +__global__ void __launch_bounds__(kThreads, 1) moeKernel(MoeParams p) +{ + extern __shared__ __align__(128) uint8_t smem[]; + uint8_t* ring = smem; + uint8_t* actRing = smem + kNS * kSlotBytes; + float* red + = reinterpret_cast(actRing + kNA * kActBytes); // kNEpi x kRed1Floats (FC1 hand-off; FC2 K-part sums) + + __shared__ __align__(8) uint64_t full[kNS], empty[kNS], actFull[kNA], actEmpty[kNA], redFull[kNEpi], + redEmpty[kNEpi], xDone, fin; + __shared__ int sList[kMaxU]; // expert of each slot (-1 = shared) + __shared__ float sWtab[kMaxU * kMaxM]; + __shared__ int sIds[kMaxM * kTopK + 32]; // + a dummy slot per lane + __shared__ float sWts[kMaxM * kTopK + 32]; + __shared__ int sNumSlots; + __shared__ unsigned sParity; + + int const tid = threadIdx.x, warp = tid / 32, lane = tid % 32, g = lane / 4, t4 = lane % 4; + int const numRows = p.numRows, numCtas = gridDim.x, c = blockIdx.x; + + if (tid == kProducerWarp * 32) + { + for (int i = 0; i < kNS; i++) + { + mbarInit(&full[i], 1); + mbarInit(&empty[i], kCWarps); + } + for (int i = 0; i < kNA; i++) + { + mbarInit(&actFull[i], 1); // the act loader's arrive.expect_tx (one bulk copy per act) + mbarInit(&actEmpty[i], kCWarps); + } + for (int i = 0; i < kNEpi; i++) + { + mbarInit(&redFull[i], kCWarps); + mbarInit(&redEmpty[i], 1); + } + mbarInit(&xDone, 1); + mbarInit(&fin, 1 + kNEpi); // the act loader and the epilogue warps are done (FC1 hand-off buffers free) + mbarInitFence(); + } + + // ---- producer helper (producer lane 0 only) ---- + uint64_t const pol = evictFirstPolicy(); + auto issueFc1 = [&](int s, int j, int e) { // stage s: tile j (expert e, -1 = shared) + int const slot = s % kNS, t = j % kTilesPerExpert; + uint8_t const* weights = e < 0 ? p.sw31 : p.w31 + static_cast(e) * (2 * kI) * kH; + float const* scales = e < 0 ? p.ss31 : p.s31 + static_cast(e) * (2 * kI / 128) * (kH / 128); + uint8_t* dst = ring + slot * kSlotBytes; + uint64_t* bar = &full[slot]; + mbarExpectTx(bar, kSlotW + 256); + bulkG2s(dst, weights + static_cast(4 * t) * kH, 4 * kH, bar, pol); // up rows 4t .. 4t + 3 + bulkG2s(dst + 4 * kH, weights + static_cast(kI + 4 * t) * kH, 4 * kH, bar, pol); // gate rows + bulkG2s(dst + kSlotW, scales + (t / 32) * (kH / 128), 128, bar); + bulkG2s(dst + kSlotW + 128, scales + (4 + t / 32) * (kH / 128), 128, bar); + }; + + // The shared expert's tiles are global tiles 0..127: CTAs c < 128 start them before the routing, and + // prefetch the shared expert's w2 (16 KB each) into L2 for the first FC2 stage. + if (tid == kProducerWarp * 32 && c < kTilesPerExpert) + { + issueFc1(0, c, -1); + l2Prefetch(p.sw2 + static_cast(c) * (kH * kI / kTilesPerExpert), kH * kI / kTilesPerExpert); + } + + cudaGridDependencySynchronize(); // x, logits, bias are produced by the previous kernel + + // ---- routing, part 1: warp t < M takes token row t (logits loads first) ---- + float sc[9]; + uint32_t key[9]; + if (warp < numRows) + { + float lg[9], bs[9]; +#pragma unroll + for (int i = 0; i < 9; i++) + { + lg[i] = __ldg(p.logits + warp * kE + lane + 32 * i); + bs[i] = __ldg(p.bias + lane + 32 * i); + } + // The score is the trtllm-gen routing kernels' sigmoid_accurate (RoutingKernel.cuh), the key its sum with + // the bias as in their SigmoidBias preprocess: the steps above 4 rows route through those kernels, and a + // different sigmoid could swap two experts within an ulp of each other between the two paths. +#pragma unroll + for (int i = 0; i < 9; i++) + { + sc[i] = 0.5f * tanhf(0.5f * lg[i]) + 0.5f; + key[i] = f2ord(sc[i] + bs[i]); + } + } + + // x -> the (still unused) act ring, for the consumers' B fragments + __shared__ __align__(8) uint64_t xBar; + if (tid == kProducerWarp * 32) + { + mbarInit(&xBar, 1); + mbarInitFence(); + mbarExpectTx(&xBar, numRows * kH * 2); + bulkG2s(actRing, p.x, numRows * kH * 2, &xBar); + } + for (int i = tid; i < kMaxU * kMaxM; i += kThreads) + { + sWtab[i] = 0.f; + } + + if (warp < numRows) + { + // top-8 of the keys (ties to the lower expert): 8 rounds of a warp max over the lanes' best keys; the winner + // lane records its expert and score in slot k and replaces its best with the best of its keys not taken yet. + // That next best (a depth-4 tree over the 9 keys) is computed by every lane while the round's reduction is in + // flight, so a round's critical path is redux -> ballot -> select. + unsigned taken = 0u; // bit i: key i is in the top-8 + uint32_t k1; + int i1; + float s1; + bestOf9(key, sc, 0u, k1, i1, s1); +#pragma unroll + for (int k = 0; k < kTopK; k++) + { + uint32_t m; + // (taken passes through the asm so that the next-best tree below is scheduled after the reduction) + asm volatile("redux.sync.max.u32 %0, %2, 0xffffffff;" : "=r"(m), "+r"(taken) : "r"(k1)); + uint32_t nk; + int ni; + float ns; + bestOf9(key, sc, taken | (1u << i1), nk, ni, ns); + asm volatile("" : "+r"(nk), "+r"(ni), "+f"(ns)); // every lane, outside the winner's path + unsigned const tied = __ballot_sync(0xffffffffu, k1 == m); + bool win = k1 == m; + if (tied & (tied - 1)) + { // several lanes hold the max: the lowest expert wins + uint32_t const cand = win ? static_cast(lane + 32 * i1) : 0xffffffffu; + win = cand == __reduce_min_sync(0xffffffffu, cand); + } + { // the losers write a private dummy slot (no branch) + int const o = win ? warp * kTopK + k : kMaxM * kTopK + lane; + sIds[o] = lane + 32 * i1; + sWts[o] = s1; + } + // row 0's picks 0, 1 (expert slots 1, 2), then 2 .. 4 (slots 3 .. 5): release the producer (one-shot + // barriers) + if (warp == 0 && k == 1) + { + __syncwarp(); + asm volatile("bar.arrive 3, 64;" ::: "memory"); + } + if (warp == 0 && k == kEarlyPicks - 1) + { + __syncwarp(); + asm volatile("bar.arrive 4, 64;" ::: "memory"); + } + taken = win ? taken | (1u << i1) : taken; + k1 = win ? nk : k1; + i1 = win ? ni : i1; + s1 = win ? ns : s1; + } + } + if (warp == kProducerWarp) + { + // The CTA's first stages after the routing are its tiles c + s G, s < kEarly: expert slots <= 5, i.e. row 0's + // first picks (slot u = row 0's pick u - 1). Issue them as soon as row 0 is done, before the other rows, + // the expert numbering and the block barriers. + int s = c < kTilesPerExpert ? 1 : 0; + asm volatile("bar.sync 3, 64;" ::: "memory"); // with warp 0: sIds[0, 1] are written + if (lane == 0) + { + for (; s < kEarly && c + s * numCtas < 3 * kTilesPerExpert; s++) + { + int const j = c + s * numCtas, u = j / kTilesPerExpert; + issueFc1(s, j, u == 0 ? -1 : sIds[u - 1]); + } + } + asm volatile("bar.sync 4, 64;" ::: "memory"); // sIds[0 .. kEarlyPicks - 1] + if (lane == 0) + { + for (; s < kEarly; s++) + { + int const j = c + s * numCtas, u = j / kTilesPerExpert; + issueFc1(s, j, u == 0 ? -1 : sIds[u - 1]); + } + } + } + else + { + if (warp == kActWarp && lane == 0) + { // this call's counter set; the CTA's next call uses the other one + unsigned const e = p.epoch[c]; + p.epoch[c] = e + 1u; + sParity = e & 1u; + } + // barriers initialized, sWtab zeroed, sIds / sWts written (every warp but the producer's) + asm volatile("bar.sync 2, %0;" ::"n"(kThreads - 32) : "memory"); + } + // ---- routing, part 2 (warp 0): distinct experts in first-occurrence order, slot 0 = the shared expert ---- + if (warp == 0) + { + int const np = numRows * kTopK; // <= 32: lane = pair (token lane / 8, rank lane % 8) + int const id = lane < np ? sIds[lane] : -2 - lane; + float w = lane < np ? sWts[lane] : 0.f; + float tot = w; + tot += __shfl_xor_sync(0xffffffffu, tot, 1); + tot += __shfl_xor_sync(0xffffffffu, tot, 2); + tot += __shfl_xor_sync(0xffffffffu, tot, 4); + unsigned const same = __match_any_sync(0xffffffffu, id); + int const first = __ffs(same) - 1; + bool const isFirst = first == lane && lane < np; + unsigned const firsts = __ballot_sync(0xffffffffu, isFirst); + int const slot = 1 + __popc(firsts & ((1u << first) - 1u)); + if (lane < np) + { + sWtab[slot * kMaxM + lane / kTopK] = w / tot * p.scaling; + if (isFirst) + { + sList[slot] = id; + } + } + if (lane == 0) + { + sList[0] = -1; + sNumSlots = 1 + __popc(firsts); + } + if (lane < numRows) + { + sWtab[lane] = 1.f; // the shared expert, weight 1 + } + } + __syncthreads(); + + // ---- the CTA's schedule ---- + // FC2: CTAs c < kFc2Ctas own 32 rows each (whole MMA tiles). The other nf CTAs only run FC1: besides their + // round-robin share of tiles [0, tl) they take the tail [tl, T1) (about the bytes of the FC2 work they skip, + // less a margin so that the last experts' activations are out before FC2 needs them). + int const numSlots = sNumSlots; + int* const cnt = p.cnt + sParity * (kMaxU * kCntPad); + int const numFc1Tiles = numSlots * kTilesPerExpert; + int const nf = numCtas - kFc2Ctas; + int const tl = max( + numFc1Tiles - max(nf * numSlots / 2 - kTailMargin, 0), kTilesPerExpert); // tiles < 128 go out before routing + int const nm = c < tl ? (tl - 1 - c) / numCtas + 1 : 0; + int const ft = c - kFc2Ctas; // rank among the FC1-only CTAs + int const nt = (ft >= 0 && tl + ft < numFc1Tiles) ? (numFc1Tiles - 1 - tl - ft) / nf + 1 : 0; + int const n1 = nm + nt; + auto tileOf = [&](int k) { return k < nm ? c + k * numCtas : tl + ft + (k - nm) * nf; }; + int const row0 = c < kFc2Ctas ? 32 * c : 0; + int const n2 = c < kFc2Ctas ? 32 : 0; + int const numFc1Stages = n1; + int const numFc2Stages = n2 > 0 ? (numSlots + 1) / 2 : 0; + int const rb0 = row0 >> 7; + int const nsc = n2 > 0 ? 1 : 0; + + if (warp == kProducerWarp) + { + // ===== producer: the CTA's FC1 stages, then its FC2 stages (two experts each) ===== + if (lane == 0) + { + int const s0 = c < kTilesPerExpert ? 1 : 0; +#pragma unroll 1 + for (int s = max(s0, kEarly); s < numFc1Stages + numFc2Stages; s++) + { + int const slot = s % kNS; + if (s >= kNS) + { + mbarWait(&empty[slot], ((s / kNS) - 1) & 1); + } + if (s < numFc1Stages) + { + int const j = tileOf(s); + issueFc1(s, j, sList[j / kTilesPerExpert]); + } + else + { + int const pr = s - numFc1Stages; + uint8_t* dst = ring + slot * kSlotBytes; + uint64_t* bar = &full[slot]; + int const ne = (2 * pr + 1 < numSlots) ? 2 : 1; + mbarExpectTx(bar, ne * (n2 * 512 + nsc * 16)); +#pragma unroll 1 + for (int i = 0; i < ne; i++) + { + int const e = sList[2 * pr + i]; + uint8_t const* weights = (e < 0 ? p.sw2 : p.w2 + static_cast(e) * kH * kI) + + static_cast(row0) * kI; + float const* scales = (e < 0 ? p.ss2 : p.s2 + static_cast(e) * (kH / 128) * (kI / 128)) + + rb0 * (kI / 128); + bulkG2s(dst + i * kFc2Pitch, weights, n2 * 512, bar, pol); + bulkG2s(dst + kSlotW + i * 32, scales, nsc * 16, bar); + } + } + } + } + __syncwarp(); + } + else if (warp == kActWarp) + { + // the other counter set (used by the previous call, used again by the next one) back to zero + if (c == 0) + { + for (int i = lane; i < kMaxU; i += 32) + { + p.cnt[(sParity ^ 1u) * (kMaxU * kCntPad) + i * kCntPad] = 0; + } + } + // ===== act loader: lane l polls expert slot l's FC1 counter (lane 0 also slot 32) ===== + if (n2 > 0) + { + uint32_t ready = 0u; + bool ready32 = false; +#pragma unroll 1 + for (int u = 0; u < numSlots; u++) + { +#pragma unroll 1 + while (!(u < 32 ? ((ready >> u) & 1u) : ready32)) + { + bool const mine = lane < numSlots && !((ready >> lane) & 1u); + bool const ok = mine && ldAcquire(cnt + lane * kCntPad) >= kTilesPerExpert; + ready |= __ballot_sync(0xffffffffu, ok); + if (numSlots > 32 && !ready32) + { + ready32 + = __shfl_sync(0xffffffffu, + static_cast(lane == 0 && ldAcquire(cnt + 32 * kCntPad) >= kTilesPerExpert), 0) + != 0; + } + __syncwarp(); + if (!(u < 32 ? ((ready >> u) & 1u) : ready32)) + { + __nanosleep(256); + } + } + // the whole [hi, lo][kMaxM][I] block (rows >= M only feed discarded MMA columns) + int const as = u % kNA; + if (u == 0) + { + mbarWait(&xDone, 0); + } + if (u >= kNA) + { + mbarWait(&actEmpty[as], ((u / kNA) - 1) & 1); + } + // one bulk copy: the lane that saw the counter complete acquired it, __syncwarp orders that before lane + // 0; the generic-proxy act stores of the other CTAs' epilogues are then made visible to the async proxy + if (lane == 0) + { + fenceProxyAsyncGlobal(); + mbarExpectTx(&actFull[as], kActBytes); + bulkG2s(actRing + as * kActBytes, p.act + static_cast(u) * 2 * kActPlane, kActBytes, + &actFull[as]); + } + __syncwarp(); + } + } + if (lane == 0) + { + mbarArrive(&fin); + } + } + else if (warp >= kEpiWarp) + { + // ===== FC1 epilogue: sum the warps' K slices, clamped SwiGLU, times the routing weight, publish ===== + int const ew = warp - kEpiWarp; +#pragma unroll 1 + for (int k = ew; k < n1; k += kNEpi) + { + int const j = tileOf(k), u = j / kTilesPerExpert, t = j % kTilesPerExpert; + int const rb = k % kNEpi; + mbarWait(&redFull[rb], (k / kNEpi) & 1); + float const* rd = red + rb * kRed1Floats; + // lane (g < 4, t4 < 2): output row 4 t + g, tokens 2 t4, 2 t4 + 1; up = MMA row g, gate = MMA row g + 4 + float up[2] = {0.f, 0.f}, gate[2] = {0.f, 0.f}; + bool const mine = g < 4 && t4 < 2; + if (mine) + { +#pragma unroll + for (int w = 0; w < kCWarps; w++) + { + float2 const vu = *reinterpret_cast(rd + (w * 16 + g * 2 + t4) * 2); + float2 const vg = *reinterpret_cast(rd + (w * 16 + (g + 4) * 2 + t4) * 2); + up[0] += vu.x; + up[1] += vu.y; + gate[0] += vg.x; + gate[1] += vg.y; + } + } + __syncwarp(); + if (lane == 0) + { + mbarArrive(&redEmpty[rb]); + } + if (mine) + { +#pragma unroll + for (int cc = 0; cc < 2; cc++) + { + int const tok = 2 * t4 + cc; + if (tok < numRows) + { + float const gv = fminf(gate[cc], p.limit); + float const uu = fminf(fmaxf(up[cc], -p.limit), p.limit); + float const act = gv / (1.f + expf(-gv)) * uu * sWtab[u * kMaxM + tok]; + __half const hi = __float2half_rn(act); + __half* dst = p.act + static_cast(u) * 2 * kActPlane + + static_cast(tok) * kActPitch + 4 * t + g; + dst[0] = hi; + dst[kActPlane] = __float2half_rn(act - __half2float(hi)); + } + } + } + __syncwarp(); + if (lane == 0) + { + redReleaseAdd(cnt + u * kCntPad, 1); + } + } + if (lane == 0) + { + mbarArrive(&fin); + } + } + else + { + // ===== consumers ===== + // B column g of the FC1 MMA: token g & 3 at K half g >> 2 (A rows 0..7: the tile's rows at K half 0, rows + // 8..15: the same rows at K half 1; only the matching (row half, column half) quadrants are used). + // This lane's fragments: x[g & 3][2048 (g >> 2) + 256 warp + 64 kb + 16 t4 + 4 j ..] (bf16 -> f16). + uint32_t xf[4][4][2]; + mbarWait(&xBar, 0); + { + __nv_bfloat16 const* xs = reinterpret_cast<__nv_bfloat16 const*>(actRing); + int const tt = g & 3, hx = g >> 2; +#pragma unroll + for (int kb = 0; kb < 4; kb++) + { +#pragma unroll + for (int j = 0; j < 4; j++) + { + uint2 raw = make_uint2(0u, 0u); + if (tt < numRows) + { + raw = *reinterpret_cast( + xs + tt * kH + 2048 * hx + 256 * warp + 64 * kb + 16 * t4 + 4 * j); + } + __half2 const h0 + = __float22half2_rn(__bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw.x))); + __half2 const h1 + = __float22half2_rn(__bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw.y))); + uint32_t const u0 = *reinterpret_cast(&h0), + u1 = *reinterpret_cast(&h1); + xf[kb][j][0] = u0; + xf[kb][j][1] = u1; + } + } + } + consumerSync(); // every consumer has its x fragments: the act ring may be overwritten + if (tid == 0) + { + mbarArrive(&xDone); + } + int s = 0; + // ---- FC1: a tile = 4 up + 4 gate rows (slot rows 0..7) over all of H; warp = K [256 warp, +256) of each half + // ---- +#pragma unroll 1 + for (int k = 0; k < n1; k++, s++) + { + int const slot = s % kNS; + mbarWait(&full[slot], (s / kNS) & 1); + uint8_t const* base = ring + slot * kSlotBytes; + float const* sc1 = reinterpret_cast(base + kSlotW) + (g < 4 ? 0 : 32); + uint8_t const* wrow = base + g * kH + 256 * warp + 16 * t4; + uint4 wl[4], wh[4]; +#pragma unroll + for (int kb = 0; kb < 4; kb++) + { + wl[kb] = *reinterpret_cast(wrow + 64 * kb); + wh[kb] = *reinterpret_cast(wrow + 2048 + 64 * kb); + } + float sl[2], sh[2]; +#pragma unroll + for (int kc = 0; kc < 2; kc++) + { + sl[kc] = sc1[2 * warp + kc]; + sh[kc] = sc1[16 + 2 * warp + kc]; + } + __syncwarp(); + if (lane == 0) + { + mbarArrive(&empty[slot]); + } + float blk[4][4]; +#pragma unroll + for (int kb = 0; kb < 4; kb++) + { +#pragma unroll + for (int i = 0; i < 4; i++) + { + blk[kb][i] = 0.f; + } + } +#pragma unroll + for (int jj = 0; jj < 4; jj++) + { +#pragma unroll + for (int kb = 0; kb < 4; kb++) + { + uint32_t const lo = word(wl[kb], jj), hi = word(wh[kb], jj); + mma16816(blk[kb], e4m3x2ToF16x2(lo & 0xffffu), e4m3x2ToF16x2(hi & 0xffffu), e4m3x2ToF16x2(lo >> 16), + e4m3x2ToF16x2(hi >> 16), xf[kb][jj][0], xf[kb][jj][1]); + } + } + // lane t4 < 2: row g x tokens 2 t4 + {0, 1} at K half 0 (c0, c1); t4 >= 2: K half 1 (c2, c3), tokens 2 t4 - + // 4 + {0, 1} + float acc[2] = {0.f, 0.f}; +#pragma unroll + for (int kc = 0; kc < 2; kc++) + { + float const s0 = t4 < 2 ? sl[kc] : sh[kc]; + float const v0 = t4 < 2 ? blk[2 * kc][0] + blk[2 * kc + 1][0] : blk[2 * kc][2] + blk[2 * kc + 1][2]; + float const v1 = t4 < 2 ? blk[2 * kc][1] + blk[2 * kc + 1][1] : blk[2 * kc][3] + blk[2 * kc + 1][3]; + acc[0] = fmaf(v0, s0, acc[0]); + acc[1] = fmaf(v1, s0, acc[1]); + } + acc[0] += __shfl_xor_sync(0xffffffffu, acc[0], 2); + acc[1] += __shfl_xor_sync(0xffffffffu, acc[1], 2); + // hand the tile's partial sums to the epilogue warp: red[warp][g * 2 + t4][2] + int const rb = k % kNEpi; + if (k >= kNEpi) + { + mbarWait(&redEmpty[rb], ((k / kNEpi) - 1) & 1); + } + float* rd = red + rb * kRed1Floats; + if (t4 < 2) + { + *reinterpret_cast(rd + (warp * 16 + g * 2 + t4) * 2) = make_float2(acc[0], acc[1]); + } + __syncwarp(); + if (lane == 0) + { + mbarArrive(&redFull[rb]); + } + } + // ---- FC2: warp = K part kp = [64 warp, 64 warp + 64) of I, both m-tiles (the CTA's rows 0..15, 16..31) ---- + int const kp = warp, kq = kp >> 1; // kq: its 128-wide K block + int rbs[2][2]; // scale block of rows mt * 16 + g (+ 8), relative to rb0 +#pragma unroll + for (int mt = 0; mt < 2; mt++) + { +#pragma unroll + for (int h = 0; h < 2; h++) + { + rbs[mt][h] = min(((row0 + mt * 16 + 8 * h + g) >> 7) - rb0, max(nsc - 1, 0)); + } + } + float acc2[2][4] = {}; + // a stage = up to two experts: wait for both acts, load every fragment, release the slots, then four + // independent MMA chains (expert x m-tile) +#pragma unroll 1 + for (int pr = 0; pr < numFc2Stages; pr++, s++) + { + int const slot = s % kNS; + int const ne = (2 * pr + 1 < numSlots) ? 2 : 1; + mbarWait(&full[slot], (s / kNS) & 1); + uint8_t const* base = ring + slot * kSlotBytes; + float const* sc2 = reinterpret_cast(base + kSlotW); + uint4 wg[2][2], wg8[2][2], b0v[2], b1v[2]; + float sA[2][2], sB[2][2]; +#pragma unroll + for (int i = 0; i < 2; i++) + { + int const u = 2 * pr + i; + if (i < ne) + { + int const as = u % kNA; + mbarWait(&actFull[as], (u / kNA) & 1); + uint8_t const* wr = base + i * kFc2Pitch + g * 512 + kp * 64 + 16 * t4; +#pragma unroll + for (int mt = 0; mt < 2; mt++) + { + wg[i][mt] = *reinterpret_cast(wr + mt * 16 * 512); + wg8[i][mt] = *reinterpret_cast(wr + (mt * 16 + 8) * 512); + sA[i][mt] = sc2[i * 8 + rbs[mt][0] * 4 + kq]; + sB[i][mt] = sc2[i * 8 + rbs[mt][1] * 4 + kq]; + } + // B column g: g < 4 -> hi of token g, g >= 4 -> lo of token g - 4 (= act row g of the [hi, lo][4] + // block) + uint8_t const* ab = actRing + as * kActBytes + (g * kActPitch + kp * 64 + 16 * t4) * 2; + b0v[i] = *reinterpret_cast(ab); + b1v[i] = *reinterpret_cast(ab + 16); + } + else + { // odd U: no second expert (zero operands, zero contribution) +#pragma unroll + for (int mt = 0; mt < 2; mt++) + { + wg[i][mt] = make_uint4(0u, 0u, 0u, 0u); + wg8[i][mt] = make_uint4(0u, 0u, 0u, 0u); + sA[i][mt] = 0.f; + sB[i][mt] = 0.f; + } + b0v[i] = make_uint4(0u, 0u, 0u, 0u); + b1v[i] = make_uint4(0u, 0u, 0u, 0u); + } + } + __syncwarp(); + if (lane == 0) + { + mbarArrive(&empty[slot]); + mbarArrive(&actEmpty[(2 * pr) % kNA]); + if (ne == 2) + { + mbarArrive(&actEmpty[(2 * pr + 1) % kNA]); + } + } + float bh[2][2][4] = {}; +#pragma unroll + for (int jj = 0; jj < 4; jj++) + { +#pragma unroll + for (int i = 0; i < 2; i++) + { + uint4 const& bv = jj < 2 ? b0v[i] : b1v[i]; + uint32_t const bb0 = (jj & 1) ? bv.z : bv.x, bb1 = (jj & 1) ? bv.w : bv.y; +#pragma unroll + for (int mt = 0; mt < 2; mt++) + { + uint32_t const aw = word(wg[i][mt], jj), aw8 = word(wg8[i][mt], jj); + mma16816(bh[i][mt], e4m3x2ToF16x2(aw & 0xffffu), e4m3x2ToF16x2(aw8 & 0xffffu), + e4m3x2ToF16x2(aw >> 16), e4m3x2ToF16x2(aw8 >> 16), bb0, bb1); + } + } + } +#pragma unroll + for (int i = 0; i < 2; i++) + { +#pragma unroll + for (int mt = 0; mt < 2; mt++) + { + acc2[mt][0] = fmaf(bh[i][mt][0], sA[i][mt], acc2[mt][0]); + acc2[mt][1] = fmaf(bh[i][mt][1], sA[i][mt], acc2[mt][1]); + acc2[mt][2] = fmaf(bh[i][mt][2], sB[i][mt], acc2[mt][2]); + acc2[mt][3] = fmaf(bh[i][mt][3], sB[i][mt], acc2[mt][3]); + } + } + } + // hi + lo: lane t4 < 2 ends up with columns 2 t4 + {0, 1} = tokens 2 t4 + {0, 1} +#pragma unroll + for (int mt = 0; mt < 2; mt++) + { +#pragma unroll + for (int x = 0; x < 4; x++) + { + acc2[mt][x] += __shfl_xor_sync(0xffffffffu, acc2[mt][x], 2); + } + } + // sum the 8 K parts (fixed order) in the FC1 hand-off buffers, free once every epilogue warp is done + mbarWait(&fin, 0); + float* rd = red; // [K part][row 0..31][token, padded to 5] + if (t4 < 2) + { +#pragma unroll + for (int mt = 0; mt < 2; mt++) + { + float* d0 = rd + (kp * 32 + mt * 16 + g) * kRed2Stride + 2 * t4; + float* d1 = d0 + 8 * kRed2Stride; + d0[0] = acc2[mt][0]; + d0[1] = acc2[mt][1]; + d1[0] = acc2[mt][2]; + d1[1] = acc2[mt][3]; + } + } + consumerSync(); + if (tid < kMaxRows2 * numRows) + { + int const row = tid % kMaxRows2, tok = tid / kMaxRows2; + float v = 0.f; +#pragma unroll + for (int q = 0; q < kCWarps; q++) + { + v += rd[(q * 32 + row) * kRed2Stride + tok]; + } + if (row < n2) + { + p.out[static_cast(tok) * kH + row0 + row] = __float2bfloat16(v); + } + } + } + cudaTriggerProgrammaticLaunchCompletion(); +} + +constexpr int kEpochSlots = 1024; // per-CTA call epochs, indexed by blockIdx.x +static_assert(kMaxCtas <= kEpochSlots, "every CTA of the largest grid has a call epoch"); + +__device__ __align__(128) __half actScratch[kMaxU * 2 * kActPlane]; +__device__ int fc1Counters[2 * kMaxU * kCntPad]; +__device__ unsigned callEpochs[kEpochSlots]; + +} // namespace + +cudaError_t decodeMoe(DecodeMoeArgs const& args, cudaStream_t stream) +{ + // Per device, set by the first call that passes the checks: the grid and the scratch symbols' addresses. + struct DeviceState + { + int grid = 0; + __half* act = nullptr; + int* cnt = nullptr; + unsigned* epoch = nullptr; + }; + + constexpr int kMaxDevices = 64; + static DeviceState states[kMaxDevices]; + int dev = 0; + cudaError_t err = cudaGetDevice(&dev); + if (err != cudaSuccess) + { + return err; + } + if (dev >= kMaxDevices) + { + return cudaErrorNotSupported; + } + DeviceState& st = states[dev]; + if (st.grid == 0) + { + int numSms = 0; + int ctasPerSm = 0; + void* act = nullptr; + void* cnt = nullptr; + void* epoch = nullptr; + err = cudaDeviceGetAttribute(&numSms, cudaDevAttrMultiProcessorCount, dev); + if (err == cudaSuccess) + { + err = cudaFuncSetAttribute(moeKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmem); + } + if (err == cudaSuccess) + { + err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(&ctasPerSm, moeKernel, kThreads, kSmem); + } + if (err == cudaSuccess) + { + err = cudaGetSymbolAddress(&act, actScratch); + } + if (err == cudaSuccess) + { + err = cudaGetSymbolAddress(&cnt, fc1Counters); + } + if (err == cudaSuccess) + { + err = cudaGetSymbolAddress(&epoch, callEpochs); + } + if (err != cudaSuccess) + { + return err; + } + // Precondition: every CTA of the grid is resident at once. FC2 CTAs spin on counters that the FC1 epilogues + // of other CTAs fill, so a CTA left waiting for an SM would hang the launch rather than fail it. The grid is + // kFc2Ctas CTAs with FC2 rows plus at least kMinFc1OnlyCtas without, one per SM, at most kMaxCtas; the + // occupancy query times the SM count the runtime reports must hold it. SMs that other kernels occupy at + // launch time are not visible here: the caller must not run work that waits on this launch concurrently. + if (numSms < kFc2Ctas + kMinFc1OnlyCtas) + { + return cudaErrorNotSupported; + } + int const grid = numSms < kMaxCtas ? numSms : kMaxCtas; + if (ctasPerSm * numSms < grid) + { + return cudaErrorCooperativeLaunchTooLarge; + } + st.act = static_cast<__half*>(act); + st.cnt = static_cast(cnt); + st.epoch = static_cast(epoch); + st.grid = grid; + } + if (args.rows < 1 || args.rows > kMaxM) + { + return cudaErrorInvalidValue; + } + MoeParams p{static_cast<__nv_bfloat16 const*>(args.x), args.logits, args.bias, args.w31, args.s31, args.sw31, + args.ss31, args.w2, args.s2, args.sw2, args.ss2, st.act, st.cnt, st.epoch, + static_cast<__nv_bfloat16*>(args.out), args.scaling, args.limit, args.rows}; + cudaLaunchConfig_t cfg{}; + cfg.gridDim = dim3(st.grid); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = kSmem; + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = args.pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, moeKernel, p); +} + +} // namespace kernels + +TRTLLM_MV2_NAMESPACE_END diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/csrc/decodeMoe.h b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/csrc/decodeMoe.h new file mode 100644 index 000000000000..910393e11e4c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/csrc/decodeMoe.h @@ -0,0 +1,75 @@ +/* + * 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. + */ +// The MoE branch of a decode step of at most 4 rows (GLM-5.3-Flash at moe_tp4), routing included. +#pragma once + +#include "tensorrt_llm/modeling_v2/common/namespace.h" + +#include + +#include + +TRTLLM_MV2_NAMESPACE_BEGIN + +namespace kernels +{ + +//! The most rows (tokens) one call takes. +constexpr int kDecodeMoeMaxRows = 4; +//! The hidden size H the kernel is compiled for. +constexpr int kDecodeMoeHidden = 4096; +//! The per-rank expert intermediate size I the kernel is compiled for (2048 at moe_tp4). +constexpr int kDecodeMoeInter = 512; +//! The number of routed experts E the kernel is compiled for. +constexpr int kDecodeMoeExperts = 288; +//! The block size of the FP8 weights' scales, along both dimensions. +constexpr int kDecodeMoeScaleBlock = 128; + +//! The operands of one decodeMoe() call. Every tensor is contiguous; the weights, their scales and x start on +//! 16-byte boundaries (they are read by TMA bulk copies). +struct DecodeMoeArgs +{ + void const* x; //!< [M, 4096] bf16 + float const* logits; //!< [M, 288] the router's outputs + float const* bias; //!< [288] the routing correction bias + uint8_t const* w31; //!< [288, 1024, 4096] e4m3: 512 up rows, then 512 gate rows + float const* s31; //!< [288, 8, 32] + uint8_t const* w2; //!< [288, 4096, 512] e4m3 + float const* s2; //!< [288, 32, 4] + uint8_t const* sw31; //!< [1024, 4096] e4m3: the shared expert, in a routed expert's layout + float const* ss31; //!< [8, 32] + uint8_t const* sw2; //!< [4096, 512] e4m3 + float const* ss2; //!< [32, 4] + void* out; //!< [M, 4096] bf16 + float scaling; //!< the routed scaling factor + float limit; //!< the SwiGLU clamp + int rows; //!< M, 1 to kDecodeMoeMaxRows + bool pdl; //!< programmatic dependent launch (only the weights are read before the grid-dependency wait) +}; + +//! One persistent launch on `stream`: routing, FC1 + SwiGLU, FC2. Its scratch (activations, per-expert completion +//! counters double-buffered by call parity, per-CTA call epochs) is device-global, so calls on one device must be +//! stream-ordered. The kernel needs every CTA resident at once (one per SM). +//! +//! Returns cudaErrorInvalidValue for a row count out of range, cudaErrorNotSupported on a device with fewer SMs than +//! the kernel's smallest grid, cudaErrorCooperativeLaunchTooLarge when the device cannot hold the whole grid at once +//! (the launch would hang rather than fail), and otherwise the result of the CUDA calls. +cudaError_t decodeMoe(DecodeMoeArgs const& args, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_MV2_NAMESPACE_END diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/csrc/glm53Fp8Ops.cpp b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/csrc/glm53Fp8Ops.cpp new file mode 100644 index 000000000000..513047925b76 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/csrc/glm53Fp8Ops.cpp @@ -0,0 +1,196 @@ +/* + * 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. + */ +// Torch ops of the target's CUDA kernels: tensors in, outputs written in place, launches on the current stream. +// The kernels are compiled for fixed shapes and read their operands with fixed strides (some through TMA bulk copies +// or 16-byte accesses), so every operand's device, dtype, sizes, contiguity and, where it matters, alignment is +// checked on the host before the launch. +#include "glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/csrc/decodeMoe.h" +#include "glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/csrc/mhcArBoundary.h" +#include "tensorrt_llm/modeling_v2/common/namespace.h" +#include "tensorrt_llm/modeling_v2/common/torchLibrary.h" + +#include +#include +#include + +#include + +TRTLLM_MV2_NAMESPACE_BEGIN + +namespace +{ + +//! The alignment TMA bulk copies and 16-byte vector accesses need. +constexpr uintptr_t kAccessAlignment = 16; + +//! Whether an operand is read by TMA bulk copies or written in 16-byte stores, so must start 16-byte aligned. +enum class Alignment +{ + kAny, + k16Bytes, +}; + +//! Checks that `tensor`, the argument `name` of op `op`, is a contiguous `dtype` tensor of sizes `sizes` on `device`, +//! starting on a 16-byte boundary when `alignment` asks for it. +void checkOperand(char const* op, char const* name, at::Tensor const& tensor, at::Device device, at::ScalarType dtype, + at::IntArrayRef sizes, Alignment alignment) +{ + TORCH_CHECK(tensor.device() == device, op, ": ", name, " must be on ", device, ", got ", tensor.device()); + TORCH_CHECK(tensor.scalar_type() == dtype, op, ": ", name, " must be ", dtype, ", got ", tensor.scalar_type()); + TORCH_CHECK(tensor.sizes() == sizes, op, ": ", name, " must have sizes ", sizes, ", got ", tensor.sizes()); + TORCH_CHECK(tensor.is_contiguous(), op, ": ", name, " must be contiguous"); + TORCH_CHECK(alignment == Alignment::kAny || reinterpret_cast(tensor.data_ptr()) % kAccessAlignment == 0, + op, ": ", name, " must start on a 16-byte boundary"); +} + +void decodeMoe(at::Tensor const& x, at::Tensor const& logits, at::Tensor const& bias, double scaling, double limit, + at::Tensor const& w31, at::Tensor const& s31, at::Tensor const& w2, at::Tensor const& s2, at::Tensor const& sw31, + at::Tensor const& ss31, at::Tensor const& sw2, at::Tensor const& ss2, at::Tensor& out, bool pdl) +{ + char const* const op = "decode_moe"; + TORCH_CHECK(x.is_cuda() && x.dim() == 2 && x.size(0) >= 1 && x.size(0) <= kernels::kDecodeMoeMaxRows, + "decode_moe: x must be a CUDA [rows, hidden] tensor of 1 to ", kernels::kDecodeMoeMaxRows, " rows"); + int64_t const rows = x.size(0); + int64_t const hidden = kernels::kDecodeMoeHidden; + int64_t const inter = kernels::kDecodeMoeInter; + int64_t const experts = kernels::kDecodeMoeExperts; + int64_t const block = kernels::kDecodeMoeScaleBlock; + at::Device const device = x.device(); + auto const fp8 = at::kFloat8_e4m3fn; + checkOperand(op, "x", x, device, at::kBFloat16, {rows, hidden}, Alignment::k16Bytes); + checkOperand(op, "logits", logits, device, at::kFloat, {rows, experts}, Alignment::kAny); + checkOperand(op, "bias", bias, device, at::kFloat, {experts}, Alignment::kAny); + checkOperand(op, "w31", w31, device, fp8, {experts, 2 * inter, hidden}, Alignment::k16Bytes); + checkOperand(op, "s31", s31, device, at::kFloat, {experts, 2 * inter / block, hidden / block}, Alignment::k16Bytes); + checkOperand(op, "w2", w2, device, fp8, {experts, hidden, inter}, Alignment::k16Bytes); + checkOperand(op, "s2", s2, device, at::kFloat, {experts, hidden / block, inter / block}, Alignment::k16Bytes); + checkOperand(op, "sw31", sw31, device, fp8, {2 * inter, hidden}, Alignment::k16Bytes); + checkOperand(op, "ss31", ss31, device, at::kFloat, {2 * inter / block, hidden / block}, Alignment::k16Bytes); + checkOperand(op, "sw2", sw2, device, fp8, {hidden, inter}, Alignment::k16Bytes); + checkOperand(op, "ss2", ss2, device, at::kFloat, {hidden / block, inter / block}, Alignment::k16Bytes); + checkOperand(op, "out", out, device, at::kBFloat16, {rows, hidden}, Alignment::kAny); + + c10::cuda::CUDAGuard const deviceGuard{device}; + kernels::DecodeMoeArgs args{}; + args.x = x.data_ptr(); + args.logits = logits.data_ptr(); + args.bias = bias.data_ptr(); + args.w31 = static_cast(w31.data_ptr()); + args.s31 = s31.data_ptr(); + args.w2 = static_cast(w2.data_ptr()); + args.s2 = s2.data_ptr(); + args.sw31 = static_cast(sw31.data_ptr()); + args.ss31 = ss31.data_ptr(); + args.sw2 = static_cast(sw2.data_ptr()); + args.ss2 = ss2.data_ptr(); + args.out = out.data_ptr(); + args.scaling = static_cast(scaling); + args.limit = static_cast(limit); + args.rows = static_cast(rows); + args.pdl = pdl; + cudaError_t const err = kernels::decodeMoe(args, at::cuda::getCurrentCUDAStream(device.index())); + TORCH_CHECK(err != cudaErrorCooperativeLaunchTooLarge, + "decode_moe: the kernel needs all of its CTAs resident at once (one per SM), and this device cannot hold its " + "grid; it would hang rather than fail"); + TORCH_CHECK(err == cudaSuccess, "decode_moe launch failed: ", cudaGetErrorString(err)); +} + +void mhcArBoundary(at::Tensor const& partial, at::Tensor const& workspace, int64_t rank, at::Tensor const& residual, + at::Tensor const& post, at::Tensor const& comb, at::Tensor const& fn, at::Tensor const& scale, + at::Tensor const& base, at::Tensor const& normWeight, int64_t sinkhornIters, double normEps, double hcEps, + double sinkhornEps, double postMult, double layerInputEps, at::Tensor& residualOut, at::Tensor& postOut, + at::Tensor& combOut, at::Tensor& layerInput, bool pdl) +{ + char const* const op = "mhc_ar_boundary"; + TORCH_CHECK(partial.is_cuda() && partial.dim() == 2 && partial.size(0) >= 1 + && partial.size(0) <= kernels::kMhcArBoundaryMaxRows, + "mhc_ar_boundary: partial must be a CUDA [rows, hidden] tensor of 1 to ", kernels::kMhcArBoundaryMaxRows, + " rows"); + TORCH_CHECK(rank >= 0 && rank < kernels::kMhcArBoundaryRanks, "mhc_ar_boundary: rank must be in [0, ", + kernels::kMhcArBoundaryRanks, "), got ", rank); + TORCH_CHECK(sinkhornIters >= 1, "mhc_ar_boundary: sinkhorn_iters must be at least 1, got ", sinkhornIters); + int64_t const rows = partial.size(0); + int64_t const hidden = kernels::kMhcArBoundaryHidden; + int64_t const mult = kernels::kMhcArBoundaryMult; + int64_t const mix = kernels::kMhcArBoundaryMix; + int64_t const numScales = 3; // the mixer's scales of pre, post and comb + at::Device const device = partial.device(); + // The workspace is the pointer table of the group's all-reduce workspace, one int64 per pointer. + checkOperand( + op, "workspace", workspace, device, at::kLong, {kernels::kMhcArBoundaryWorkspacePointers}, Alignment::kAny); + checkOperand(op, "partial", partial, device, at::kBFloat16, {rows, hidden}, Alignment::k16Bytes); + checkOperand(op, "residual", residual, device, at::kBFloat16, {rows, mult, hidden}, Alignment::k16Bytes); + checkOperand(op, "post", post, device, at::kFloat, {rows, mult}, Alignment::k16Bytes); + checkOperand(op, "comb", comb, device, at::kFloat, {rows, mult, mult}, Alignment::k16Bytes); + checkOperand(op, "fn", fn, device, at::kFloat, {mix, mult * hidden}, Alignment::k16Bytes); + checkOperand(op, "scale", scale, device, at::kFloat, {numScales}, Alignment::kAny); + checkOperand(op, "base", base, device, at::kFloat, {mix}, Alignment::kAny); + checkOperand(op, "norm_weight", normWeight, device, at::kBFloat16, {hidden}, Alignment::kAny); + checkOperand(op, "residual_out", residualOut, device, at::kBFloat16, {rows, mult, hidden}, Alignment::k16Bytes); + checkOperand(op, "post_out", postOut, device, at::kFloat, {rows, mult}, Alignment::kAny); + checkOperand(op, "comb_out", combOut, device, at::kFloat, {rows, mult, mult}, Alignment::kAny); + checkOperand(op, "layer_input", layerInput, device, at::kBFloat16, {rows, hidden}, Alignment::k16Bytes); + + c10::cuda::CUDAGuard const deviceGuard{device}; + kernels::MhcArBoundaryArgs args{}; + args.workspace = static_cast(workspace.data_ptr()); + args.rank = static_cast(rank); + args.partial = partial.data_ptr(); + args.residual = residual.data_ptr(); + args.postPrev = post.data_ptr(); + args.combPrev = comb.data_ptr(); + args.fn = fn.data_ptr(); + args.scale = scale.data_ptr(); + args.base = base.data_ptr(); + args.normWeight = normWeight.data_ptr(); + args.residualOut = residualOut.data_ptr(); + args.postOut = postOut.data_ptr(); + args.combOut = combOut.data_ptr(); + args.layerInput = layerInput.data_ptr(); + args.rows = static_cast(rows); + args.sinkhornIters = static_cast(sinkhornIters); + args.normEps = static_cast(normEps); + args.hcEps = static_cast(hcEps); + args.sinkhornEps = static_cast(sinkhornEps); + args.postMult = static_cast(postMult); + args.layerInputEps = static_cast(layerInputEps); + args.pdl = pdl; + cudaError_t const err = kernels::mhcArBoundary(args, at::cuda::getCurrentCUDAStream(device.index())); + TORCH_CHECK(err == cudaSuccess, "mhc_ar_boundary launch failed: ", cudaGetErrorString(err)); +} + +} // namespace + +TRTLLM_MV2_NAMESPACE_END + +TRTLLM_MV2_LIBRARY_FRAGMENT(m) +{ + m.def( + "decode_moe(Tensor x, Tensor logits, Tensor bias, float scaling, float limit, Tensor w31, Tensor s31, " + "Tensor w2, Tensor s2, Tensor sw31, Tensor ss31, Tensor sw2, Tensor ss2, Tensor(a!) out, bool pdl) -> ()"); + m.def( + "mhc_ar_boundary(Tensor partial, Tensor workspace, int rank, Tensor residual, Tensor post, Tensor comb, " + "Tensor fn, Tensor scale, Tensor base, Tensor norm_weight, int sinkhorn_iters, float norm_eps, " + "float hc_eps, float sinkhorn_eps, float post_mult, float layer_input_eps, Tensor(a!) residual_out, " + "Tensor(b!) post_out, Tensor(c!) comb_out, Tensor(d!) layer_input, bool pdl) -> ()"); +} + +TRTLLM_MV2_LIBRARY_IMPL(CUDA, m) +{ + m.impl("decode_moe", &TRTLLM_MV2_NS::decodeMoe); + m.impl("mhc_ar_boundary", &TRTLLM_MV2_NS::mhcArBoundary); +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/csrc/mhcArBoundary.cu b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/csrc/mhcArBoundary.cu new file mode 100644 index 000000000000..239354f965ed --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/csrc/mhcArBoundary.cu @@ -0,0 +1,753 @@ +/* + * 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. + */ +// One-shot all-reduce of a decode step's partial rows fused with the mHC boundary that follows it (GLM-5.3-Flash at +// tp4: hidden 4096, hc_mult 4, at most 8 rows), as one cluster of 16 CTAs. +// +// All-reduce: TensorRT-LLM's Lamport one-shot protocol on the all-reduce workspace (the LamportComm layout: three +// rotating receive buffers per rank, the flag, the counter and the clear size), so calls interleave with the +// library's one-shot all-reduce on the same workspace. Each CTA pushes its 256 columns of the rows to every rank +// (-0.0 marks a word not yet arrived, so a -0.0 value is sent as +0.0), clears its share of the previous call's +// buffer, polls until the four ranks' slices are there, sums them in rank order in FP32 and rounds to BF16, like the +// library's one-shot kernel. The layout and rotation are in namespace lamport below. +// mHC boundary (post-mapping of the previous block, pre-mapping of the next one): every CTA computes the new streams +// of its columns, its partial mixer products (TF32 MMA, the streams are BF16 and exact in TF32) and Gram entries; +// the 16 CTAs exchange their sums through distributed shared memory and each adds them in CTA order +// (deterministic). Each CTA then derives the coefficients (the aux warp runs the Sinkhorn normalization) and writes +// its columns of the streams and of the layer input, RMS-normalized through the Gram matrix (no second reduction). +// The CTA's 96 KB slice of the mixer weights is loaded before the grid-dependency wait. +#include "glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/csrc/mhcArBoundary.h" + +#include +#include +#include + +#include + +TRTLLM_MV2_NAMESPACE_BEGIN + +namespace kernels +{ +namespace +{ + +namespace cg = cooperative_groups; + +constexpr int kN = kMhcArBoundaryMult, kMix = kMhcArBoundaryMix, kGram = 10, + kAcc = kMix + kGram; // 24 + 10 sums per row +constexpr int kH = kMhcArBoundaryHidden, kNH = kN * kH, kRanks = kMhcArBoundaryRanks; +constexpr int kCtas = 16, kCols = kH / kCtas, kWarps = kCols / 32, kThreads = kCols + 32; // + aux warp +constexpr int kMaxRows = kMhcArBoundaryMaxRows; +constexpr int kChunks = kCols / 8; // 16-byte chunks of a row slice +constexpr int kSlice = kCols + 8, kTok = kN * kSlice + 8; // padded (bank-conflict free MMA fragment loads) +constexpr int kFnRow = kN * kCols + 4; // fn rows (o, j) at o * kFnRow + j * kCols (floats) +constexpr int kFnBytes = kMix * kFnRow * 4; +constexpr int kResBytes = kMaxRows * kTok * 2; // this CTA's columns of the residual streams +constexpr int kSmem = kResBytes + kWarps * kMaxRows * kAcc * 4 + kMaxRows * kCols * 2 + kMaxRows * 20 * 4 + kFnBytes; + +//! The Lamport state of TensorRT-LLM's all-reduce workspace, which this kernel shares with the library's one-shot +//! all-reduce. The library keeps the layout private: LamportComm in +//! cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu (and its copy in +//! MiniMaxReduceRMSKernel.cu) reads it, and allocate_allreduce_fusion_workspace in +//! tensorrt_llm/_torch/distributed/allreduce_helper.py builds it. Every value here is that code's; the source +//! test test_modeling_v2_glm53_fp8_lamport_layout.py checks them against that code, so changing one side fails CI. +namespace lamport +{ +//! The pointer table: kRanks IPC buffers, kRanks barrier flags, kRanks Lamport buffers, then the two below. +constexpr int kBufferSlot = 2 * kRanks; //!< ws[kBufferSlot + r]: rank r's Lamport buffers (kNumBuffers of them) +constexpr int kFlagSlot = 3 * kRanks; //!< ws[kFlagSlot]: int32 words [counter, (non-Lamport flag), flag] +constexpr int kLayoutSlot = 3 * kRanks + 1; //!< ws[kLayoutSlot]: int64 words [clear size, buffer size], in bytes +constexpr int kTableSize = 3 * kRanks + 2; //!< pointers in the table +constexpr int kCounterWord = 0; //!< CTAs of the current call that have read the flag +constexpr int kFlagWord = 2; //!< the call's flag; its receive buffer is flag % kNumBuffers +constexpr int kClearSizeWord = 0; //!< bytes the previous call wrote, which this call clears +constexpr int kCommSizeWord = 1; //!< bytes of one receive buffer +constexpr int kNumBuffers = 3; //!< receive buffers per rank, rotated by the flag + +//! The receive buffer of the call with flag `flag`. +__device__ __forceinline__ int dataBuffer(int flag) +{ + return flag % kNumBuffers; +} + +//! The receive buffer of the call before the one with flag `flag`, which that call clears. +__device__ __forceinline__ int clearBuffer(int flag) +{ + return (flag + kNumBuffers - 1) % kNumBuffers; +} + +//! The flag of the call after the one with flag `flag`. +__device__ __forceinline__ int nextFlag(int flag) +{ + return (flag + 1) % kNumBuffers; +} +} // namespace lamport + +static_assert(lamport::kTableSize == kMhcArBoundaryWorkspacePointers, "the workspace pointer table the binding checks"); + +//! The kernel's arguments (MhcArBoundaryArgs, typed). +struct Params +{ + void** workspace; //!< TensorRT-LLM all-reduce workspace + int rank; //!< this rank in the group + __nv_bfloat16 const* partial; //!< [M, H] this rank's partial block output + __nv_bfloat16 const* res; //!< [M, N, H] + float const* postPrev; //!< [M, N] + float const* combPrev; //!< [M, N, N] + float const* fn; //!< [24, N * H] + float const* scale; //!< [3] + float const* base; //!< [24] + __nv_bfloat16 const* w; //!< [H] + __nv_bfloat16* resOut; //!< [M, N, H] + float* postOut; //!< [M, N] + float* combOut; //!< [M, N, N] + __nv_bfloat16* liOut; //!< [M, H] + int numRows; //!< M + int skIters; //!< Sinkhorn iterations + float normEps; //!< the mixer's RMS epsilon + float hcEps; //!< added to pre + float skEps; //!< the Sinkhorn epsilon + float postMult; //!< post's multiplier + float liEps; //!< the layer input's RMSNorm epsilon +}; + +__device__ __forceinline__ constexpr int gramIdx(int j, int k) +{ + return j * kN - j * (j - 1) / 2 + (k - j); +} + +__device__ __forceinline__ bool isNegZero(uint32_t v) +{ + return v == 0x80000000u; +} + +__device__ __forceinline__ uint4 ldVolatile(uint4 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)); + return v; +} + +__device__ __forceinline__ uint32_t toTf32(float v) +{ + uint32_t r; + asm("cvt.rna.tf32.f32 %0, %1;" : "=r"(r) : "f"(v)); + return r; +} + +// D += A (16 x 8, row) * B (8 x 8, col), TF32 inputs, FP32 accumulation +__device__ __forceinline__ void mmaTf32( + 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.m16n8k8.row.col.f32.tf32.tf32.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)); +} + +__device__ __forceinline__ uint32_t f32Bits(float v) +{ + return __float_as_uint(v); +} + +__device__ __forceinline__ uint32_t smemU32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__global__ void __cluster_dims__(kCtas, 1, 1) __launch_bounds__(kThreads, 1) mhcArKernel(Params p) +{ + extern __shared__ __align__(128) uint8_t smem[]; + __nv_bfloat16* sRes = reinterpret_cast<__nv_bfloat16*>(smem); // [rows][4][256], then the new streams + float* sWarp = reinterpret_cast(smem + kResBytes); // [warps][rows][34] + __nv_bfloat16* sX + = reinterpret_cast<__nv_bfloat16*>(smem + kResBytes + kWarps * kMaxRows * kAcc * 4); // [rows][256] + float* sPc = reinterpret_cast(smem + kResBytes + kWarps * kMaxRows * kAcc * 4 + kMaxRows * kCols * 2); + // sPc: postPrev [rows][4], then combPrev [rows][16] + float* sFn = sPc + kMaxRows * 20; + __shared__ float sAll[kCtas * kMaxRows * kAcc]; // every CTA's sums, slot c written by CTA c (DSMEM stores) + __shared__ float sTot[kMaxRows * kAcc]; // the 16 CTAs' totals + cg::cluster_group cluster = cg::this_cluster(); + __shared__ float sSk[2][kN + kN * kN]; // this CTA's rows' post and comb (aux warp), stored after the wait + __shared__ __align__(8) uint64_t inBar, fnBar; + + int const cta = blockIdx.x, t = threadIdx.x, warp = t / 32, lane = t % 32; + int const c0 = cta * kCols, h = c0 + t, numRows = p.numRows; + bool const colThread = warp < kWarps, aux = warp == kWarps; + __shared__ float sCoef[kMaxRows][kN]; // layer_input = sum_j coef[m][j] r_j * w + + // before the PDL wait: weights (this column's 96 mixer weights, kept in registers for every row, and its norm + // weight) and the all-reduce workspace's pointer table and buffer size, which never change + // mixer K slice of this warp: stream w / 2, columns [128 (w % 2), + 128), 16 k-steps of 8 (B from sFn) + int const g = lane >> 2, tig = lane & 3; + int const mj = warp >> 1, mcol = (warp & 1) * 128; + float const wh = colThread ? __bfloat162float(p.w[h]) : 0.f; + void** ws = p.workspace; + int* const flagWords = static_cast(ws[lamport::kFlagSlot]); + long long* const layoutWords = static_cast(ws[lamport::kLayoutSlot]); + int* counterPtr = flagWords + lamport::kCounterWord; + int* flagPtr = flagWords + lamport::kFlagWord; + long long* clearPtr = layoutWords + lamport::kClearSizeWord; + long long const commSize = layoutWords[lamport::kCommSizeWord]; + uint8_t* bufs[kRanks]; +#pragma unroll + for (int r = 0; r < kRanks; r++) + { + bufs[r] = static_cast(ws[lamport::kBufferSlot + r]); + } + if (t == kCols) + { + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(smemU32(&inBar))); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(smemU32(&fnBar))); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + asm volatile( + "mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(smemU32(&fnBar)), "r"(kMix * kN * kCols * 4) + : "memory"); + } + __syncthreads(); + // the fn slice (weights) before the wait: 96 rows of 1 KB + if (aux) + { + for (int i = lane; i < kMix * kN; i += 32) + { + int const o = i / kN, j = i % kN; + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( + smemU32(sFn + o * kFnRow + j * kCols)), + "l"(p.fn + static_cast(o) * kNH + j * kH + c0), "r"(kCols * 4), "r"(smemU32(&fnBar)) + : "memory"); + } + } + + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); // this CTA has started + // the producer of the partial rows (and the previous boundary, the last user of the Lamport state) has completed + cudaGridDependencySynchronize(); + + // aux warp: stage this CTA's columns of the residual streams and the rows' post / comb (overlaps the all-reduce) + if (aux) + { + if (lane == 0) + { + uint32_t const bytes = numRows * kN * kCols * 2 + numRows * 20 * 4; + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(smemU32(&inBar)), "r"(bytes) + : "memory"); + } + __syncwarp(); + for (int i = lane; i < numRows * kN + 2; i += 32) + { + bool const resCopy = i < numRows * kN; + uint32_t const dst = resCopy ? smemU32(sRes + (i / kN) * kTok + (i % kN) * kSlice) + : smemU32(sPc + (i == numRows * kN ? 0 : kMaxRows * 4)); + void const* src = resCopy + ? static_cast(p.res + static_cast(i) * kH + c0) + : (i == numRows * kN ? static_cast(p.postPrev) : static_cast(p.combPrev)); + uint32_t const bytes = resCopy ? kCols * 2 : (i == numRows * kN ? numRows * 16 : numRows * 64); + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"(dst), + "l"(src), "r"(bytes), "r"(smemU32(&inBar)) + : "memory"); + } + } + // Lamport state (TensorRT-LLM LamportComm): the receive buffers rotate with the flag, and the previous call's is + // cleared. The clear size is the bytes that call wrote, whatever kernel or dtype it was; it is cleared in whole + // 16-byte stores, rounded up. + int const flag = *flagPtr; + long long const clearSize = *clearPtr; + if (t == 0) + { + atomicAdd(counterPtr, 1); + } + uint8_t* dataBufs[kRanks]; +#pragma unroll + for (int r = 0; r < kRanks; r++) + { + dataBufs[r] = bufs[r] + lamport::dataBuffer(flag) * commSize; + } + uint4* clearBuf = reinterpret_cast(bufs[p.rank] + lamport::clearBuffer(flag) * commSize); + int const tot = numRows * kH / 8; // 16-byte accesses in one rank's message + + // push this CTA's columns of the partial rows to every rank (-0.0 words sent as +0.0) + for (int i = t; i < numRows * kChunks; i += kThreads) + { + int const m = i / kChunks, q = i % kChunks, idx = (m * kH + c0) / 8 + q; + uint4 v = reinterpret_cast(p.partial)[idx]; + v.x = isNegZero(v.x) ? 0u : v.x; + v.y = isNegZero(v.y) ? 0u : v.y; + v.z = isNegZero(v.z) ? 0u : v.z; + v.w = isNegZero(v.w) ? 0u : v.w; +#pragma unroll + for (int r = 0; r < kRanks; r++) + { + reinterpret_cast(dataBufs[r])[p.rank * tot + idx] = v; + } + } + // clear this CTA's share of the previous call's buffer + uint4 const sentinel = make_uint4(0x80000000u, 0x80000000u, 0x80000000u, 0x80000000u); + for (long long i = static_cast(cta) * kThreads + t; i < (clearSize + 15) / 16; + i += static_cast(kCtas) * kThreads) + { + clearBuf[i] = sentinel; + } + // poll and sum the four ranks' slices in rank order (FP32), rounded to BF16 like the library's all-reduce + for (int i = t; i < numRows * kChunks; i += kThreads) + { + int const m = i / kChunks, q = i % kChunks, idx = (m * kH + c0) / 8 + q; + uint4 v[kRanks]; + bool done; + do + { + done = true; +#pragma unroll + for (int r = 0; r < kRanks; r++) + { + v[r] = ldVolatile(reinterpret_cast(dataBufs[p.rank]) + r * tot + idx); + done &= !(isNegZero(v[r].x) || isNegZero(v[r].y) || isNegZero(v[r].z) || isNegZero(v[r].w)); + } + } while (!done); + float acc[8]; +#pragma unroll + for (int e = 0; e < 8; e++) + { + acc[e] = 0.f; + } +#pragma unroll + for (int r = 0; r < kRanks; r++) + { + __nv_bfloat16 const* bv = reinterpret_cast<__nv_bfloat16 const*>(&v[r]); +#pragma unroll + for (int e = 0; e < 8; e++) + { + acc[e] += __bfloat162float(bv[e]); + } + } + uint4 o; + __nv_bfloat16* ob = reinterpret_cast<__nv_bfloat16*>(&o); +#pragma unroll + for (int e = 0; e < 8; e++) + { + ob[e] = __float2bfloat16(acc[e]); + } + reinterpret_cast(sX + m * kCols)[q] = o; + } + asm volatile( + "{\n .reg .pred p;\n WAIT_%=:\n mbarrier.try_wait.parity.shared::cta.b64 p, [%0], 0;\n @!p bra WAIT_%=;\n}" :: + "r"(smemU32(&inBar)) + : "memory"); + __syncthreads(); + + // post-mapping of column h: the new streams, in place; four rows per group, the group's inputs loaded first + float4 const* sPost4 = reinterpret_cast(sPc); + float4 const* sComb4 = reinterpret_cast(sPc + kMaxRows * 4); + if (colThread) + { + for (int m0 = 0; m0 < numRows; m0 += 4) + { + float xv[4], rp[4][kN]; + float4 po[4], cm[4][kN]; +#pragma unroll + for (int k = 0; k < 4; k++) + { + int const m = min(m0 + k, numRows - 1); + xv[k] = __bfloat162float(sX[m * kCols + t]); +#pragma unroll + for (int i = 0; i < kN; i++) + { + rp[k][i] = __bfloat162float(sRes[m * kTok + i * kSlice + t]); + } + po[k] = sPost4[m]; +#pragma unroll + for (int i = 0; i < kN; i++) + { + cm[k][i] = sComb4[m * kN + i]; + } + } +#pragma unroll + for (int k = 0; k < 4; k++) + { + int const m = m0 + k; + if (m < numRows) + { + float const pj[kN] = {po[k].x, po[k].y, po[k].z, po[k].w}; +#pragma unroll + for (int j = 0; j < kN; j++) + { + float a = pj[j] * xv[k]; +#pragma unroll + for (int i = 0; i < kN; i++) + { + float const c = j == 0 ? cm[k][i].x + : j == 1 ? cm[k][i].y + : j == 2 ? cm[k][i].z + : cm[k][i].w; + a += c * rp[k][i]; + } + __nv_bfloat16 const bb = __float2bfloat16(a); + sRes[m * kTok + j * kSlice + t] = bb; + } + } + } + } + } + __syncthreads(); + // this warp's partial mixer outputs (K slice above) and Gram entries (columns [32 w, 32 w + 32)) on the + // tensor cores; A operands from the new streams (BF16, exact in TF32) + asm volatile( + "{\n .reg .pred p;\n WAIT_%=:\n mbarrier.try_wait.parity.shared::cta.b64 p, [%0], 0;\n @!p bra WAIT_%=;\n}" :: + "r"(smemU32(&fnBar)) + : "memory"); + if (colThread) + { + auto rv = [&](int m, int j, int col) -> uint32_t + { return m < numRows ? f32Bits(__bfloat162float(sRes[m * kTok + j * kSlice + col])) : 0u; }; + for (int mt = 0; mt * 16 < numRows; mt++) + { + float acc[3][4] = {}; + int const r0 = 16 * mt + g, r1 = r0 + 8; +#pragma unroll + for (int st = 0; st < 16; st++) + { + int const col = mcol + 8 * st + tig; + uint32_t const a0 = rv(r0, mj, col), a1 = rv(r1, mj, col), a2 = rv(r0, mj, col + 4), + a3 = rv(r1, mj, col + 4); +#pragma unroll + for (int nt = 0; nt < 3; nt++) + { + float const* fb = sFn + (nt * 8 + g) * kFnRow + mj * kCols + col; + mmaTf32(acc[nt], a0, a1, a2, a3, toTf32(fb[0]), toTf32(fb[4])); + } + } +#pragma unroll + for (int nt = 0; nt < 3; nt++) + { + int const n = nt * 8 + 2 * tig; + if (r0 < numRows) + { + sWarp[(warp * kMaxRows + r0) * kAcc + n] = acc[nt][0]; + sWarp[(warp * kMaxRows + r0) * kAcc + n + 1] = acc[nt][1]; + } + if (r1 < numRows) + { + sWarp[(warp * kMaxRows + r1) * kAcc + n] = acc[nt][2]; + sWarp[(warp * kMaxRows + r1) * kAcc + n + 1] = acc[nt][3]; + } + } + } + // Gram: rows (token 4 mg + r / 4, stream r % 4) x columns (token 4 mg + 2 q + n / 4, stream n % 4) + for (int mg = 0; mg * 4 < numRows; mg++) + { + float ga[2][4] = {}; + int const ta0 = 4 * mg + (g >> 2), ta1 = ta0 + 2, sa = g & 3; +#pragma unroll + for (int st = 0; st < 4; st++) + { + int const col = 32 * warp + 8 * st + tig; + uint32_t const a0 = rv(ta0, sa, col), a1 = rv(ta1, sa, col), a2 = rv(ta0, sa, col + 4), + a3 = rv(ta1, sa, col + 4); +#pragma unroll + for (int q = 0; q < 2; q++) + { + int const tb = 4 * mg + 2 * q + (g >> 2); + mmaTf32(ga[q], a0, a1, a2, a3, rv(tb, sa, col), rv(tb, sa, col + 4)); + } + } +#pragma unroll + for (int q = 0; q < 2; q++) + { + int const tb = 2 * q + (tig >> 1), kb0 = 2 * (tig & 1); +#pragma unroll + for (int c = 0; c < 4; c++) + { + int const ta = (c < 2 ? 0 : 2) + (g >> 2), kb = kb0 + (c & 1), token = 4 * mg + ta; + if (ta == tb && sa <= kb && token < numRows) + { + sWarp[(warp * kMaxRows + token) * kAcc + kMix + gramIdx(sa, kb)] = ga[q][c]; + } + } + } + } + } + __syncthreads(); + // this CTA's sums into slot `cta` of every CTA's s_all, then each CTA adds the 16 slots in CTA order + asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); // every CTA of the cluster has started + for (int i = t; i < numRows * kAcc; i += kThreads) + { + float s3 = 0.f; +#pragma unroll + for (int w8 = 0; w8 < kWarps; w8++) + { + s3 += sWarp[(w8 * kMaxRows) * kAcc + i]; + } +#pragma unroll + for (int c = 0; c < kCtas; c++) + { + cluster.map_shared_rank(sAll, c)[cta * kMaxRows * kAcc + i] = s3; + } + } + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + for (int i = t; i < numRows * kAcc; i += kThreads) + { + float s3 = 0.f; +#pragma unroll + for (int c = 0; c < kCtas; c++) + { + s3 += sAll[c * kMaxRows * kAcc + i]; + } + sTot[i] = s3; + } + __syncthreads(); + // per row: the collapse coefficients (pre_j and the RMSNorm scale of the collapsed row from the Gram matrix) + if (warp == 0 && lane < numRows) + { + float const* a = sTot + lane * kAcc; + float gram[kN][kN]; +#pragma unroll + for (int j = 0; j < kN; j++) + { +#pragma unroll + for (int i = j; i < kN; i++) + { + float const entry = a[kMix + gramIdx(j, i)]; + gram[j][i] = entry; + gram[i][j] = entry; + } + } + float const rs + = rsqrtf((gram[0][0] + gram[1][1] + gram[2][2] + gram[3][3]) / static_cast(kNH) + p.normEps); + float pre[kN]; +#pragma unroll + for (int j = 0; j < kN; j++) + { + pre[j] = 1.f / (1.f + __expf(-(a[j] * rs * p.scale[0] + p.base[j]))) + p.hcEps; + } + float q = 0.f; +#pragma unroll + for (int j = 0; j < kN; j++) + { +#pragma unroll + for (int i = 0; i < kN; i++) + { + q += pre[j] * pre[i] * gram[j][i]; + } + } + float const sc = rsqrtf(q / static_cast(kH) + p.liEps); +#pragma unroll + for (int j = 0; j < kN; j++) + { + sCoef[lane][j] = pre[j] * sc; + } + } + __syncthreads(); + if (colThread) + { + // layer_input of column h for every row +#pragma unroll + for (int m = 0; m < kMaxRows; m++) + { + if (m < numRows) + { + float y = 0.f; +#pragma unroll + for (int j = 0; j < kN; j++) + { + y += sCoef[m][j] * __bfloat162float(sRes[m * kTok + j * kSlice + t]); + } + sX[m * kCols + t] = __float2bfloat16(y * wh); // s_x (the all-reduced rows) is free again + } + } + } + else if (lane < (kMaxRows + kCtas - 1) / kCtas && cta + kCtas * lane < numRows) + { + // aux warp: post and the Sinkhorn-normalized comb of the rows m with m % 16 == cta, one lane each + int const m = cta + kCtas * lane; + float const* a = sTot + m * kAcc; + float d[kN]; +#pragma unroll + for (int j = 0; j < kN; j++) + { + d[j] = a[kMix + gramIdx(j, j)]; + } + float const rs = rsqrtf((d[0] + d[1] + d[2] + d[3]) / static_cast(kNH) + p.normEps); +#pragma unroll + for (int j = 0; j < kN; j++) + { + sSk[lane][j] = p.postMult / (1.f + __expf(-(a[kN + j] * rs * p.scale[1] + p.base[kN + j]))); + } + float c[kN][kN]; +#pragma unroll + for (int i = 0; i < kN; i++) + { + float mx = -3.4e38f; +#pragma unroll + for (int j = 0; j < kN; j++) + { + c[i][j] = a[2 * kN + i * kN + j] * rs * p.scale[2] + p.base[2 * kN + i * kN + j]; + mx = fmaxf(mx, c[i][j]); + } + float s3 = 0.f; +#pragma unroll + for (int j = 0; j < kN; j++) + { + c[i][j] = __expf(c[i][j] - mx); + s3 += c[i][j]; + } + float const inv = 1.f / s3; +#pragma unroll + for (int j = 0; j < kN; j++) + { + c[i][j] = c[i][j] * inv + p.skEps; + } + } + for (int it = 0; it < p.skIters; it++) + { + if (it > 0) + { +#pragma unroll + for (int i = 0; i < kN; i++) + { + float const inv = 1.f / (c[i][0] + c[i][1] + c[i][2] + c[i][3] + p.skEps); +#pragma unroll + for (int j = 0; j < kN; j++) + { + c[i][j] *= inv; + } + } + } +#pragma unroll + for (int j = 0; j < kN; j++) + { + float const inv = 1.f / (c[0][j] + c[1][j] + c[2][j] + c[3][j] + p.skEps); +#pragma unroll + for (int i = 0; i < kN; i++) + { + c[i][j] *= inv; + } + } + } +#pragma unroll + for (int i = 0; i < kN; i++) + { +#pragma unroll + for (int j = 0; j < kN; j++) + { + sSk[lane][kN + i * kN + j] = c[i][j]; + } + } + } + __syncthreads(); + // 16-byte stores: this CTA's 256 columns of every residual stream row and layer_input row + for (int i = t; i < numRows * (kN + 1) * kChunks; i += kThreads) + { + int const sl = i / kChunks, q = i % kChunks, m = sl / (kN + 1), j = sl % (kN + 1); + if (j < kN) + { + reinterpret_cast(p.resOut + (static_cast(m) * kN + j) * kH + c0)[q] + = reinterpret_cast(sRes + m * kTok + j * kSlice)[q]; + } + else + { + reinterpret_cast(p.liOut + static_cast(m) * kH + c0)[q] + = reinterpret_cast(sX + m * kCols)[q]; + } + } + if (!colThread) + { + // aux warp: its rows' post and comb + if (lane < (kMaxRows + kCtas - 1) / kCtas && cta + kCtas * lane < numRows) + { + int const m = cta + kCtas * lane; +#pragma unroll + for (int j = 0; j < kN; j++) + { + p.postOut[m * kN + j] = sSk[lane][j]; + } +#pragma unroll + for (int i = 0; i < kN * kN; i++) + { + p.combOut[m * kN * kN + i] = sSk[lane][kN + i]; + } + } + } + cudaTriggerProgrammaticLaunchCompletion(); // after the outputs: the next kernel's CTAs do not crowd them + // Lamport rotation (LamportComm::update): once every CTA has counted in + if (cta == 0 && t == 0) + { + while (*static_cast(counterPtr) != kCtas) + { + } + *flagPtr = lamport::nextFlag(flag); + // bytes, as LamportComm + *clearPtr = static_cast(numRows) * kH * kRanks * static_cast(sizeof(__nv_bfloat16)); + *counterPtr = 0; + } +} + +} // namespace + +cudaError_t mhcArBoundary(MhcArBoundaryArgs const& args, cudaStream_t stream) +{ + if (args.rows < 1 || args.rows > kMaxRows) + { + return cudaErrorInvalidValue; + } + constexpr int kMaxDevices = 64; + static bool attributesSet[kMaxDevices] = {}; + int dev = 0; + cudaError_t err = cudaGetDevice(&dev); + if (err != cudaSuccess) + { + return err; + } + if (dev >= kMaxDevices) + { + return cudaErrorNotSupported; + } + if (!attributesSet[dev]) + { + err = cudaFuncSetAttribute(mhcArKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmem); + if (err == cudaSuccess) + { + err = cudaFuncSetAttribute(mhcArKernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1); + } + if (err != cudaSuccess) + { + return err; + } + attributesSet[dev] = true; + } + Params p{args.workspace, args.rank, static_cast<__nv_bfloat16 const*>(args.partial), + static_cast<__nv_bfloat16 const*>(args.residual), args.postPrev, args.combPrev, args.fn, args.scale, args.base, + static_cast<__nv_bfloat16 const*>(args.normWeight), static_cast<__nv_bfloat16*>(args.residualOut), args.postOut, + args.combOut, static_cast<__nv_bfloat16*>(args.layerInput), args.rows, args.sinkhornIters, args.normEps, + args.hcEps, args.sinkhornEps, args.postMult, args.layerInputEps}; + cudaLaunchConfig_t cfg{}; + cfg.gridDim = dim3(kCtas); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = kSmem; + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = args.pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, mhcArKernel, p); +} + +} // namespace kernels + +TRTLLM_MV2_NAMESPACE_END diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/csrc/mhcArBoundary.h b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/csrc/mhcArBoundary.h new file mode 100644 index 000000000000..a722e7101fc8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/csrc/mhcArBoundary.h @@ -0,0 +1,79 @@ +/* + * 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. + */ +// The tensor-parallel all-reduce of a decode step's partial rows fused with the mHC boundary after it +// (GLM-5.3-Flash at tp4). +#pragma once + +#include "tensorrt_llm/modeling_v2/common/namespace.h" + +#include + +TRTLLM_MV2_NAMESPACE_BEGIN + +namespace kernels +{ + +//! The most rows (tokens) one call takes. +constexpr int kMhcArBoundaryMaxRows = 8; +//! The hidden size the kernel is compiled for. +constexpr int kMhcArBoundaryHidden = 4096; +//! The number of residual streams (hc_mult) the kernel is compiled for. +constexpr int kMhcArBoundaryMult = 4; +//! The size of the tensor-parallel group the kernel is compiled for. +constexpr int kMhcArBoundaryRanks = 4; +//! The mixer's outputs: pre and post (one per stream) and comb (one per pair of streams). +constexpr int kMhcArBoundaryMix = (2 + kMhcArBoundaryMult) * kMhcArBoundaryMult; +//! The pointers in the all-reduce workspace's table: three buffer sets per rank, the flag and the layout buffers. +constexpr int kMhcArBoundaryWorkspacePointers = 3 * kMhcArBoundaryRanks + 2; + +//! The operands of one mhcArBoundary() call. Every tensor is contiguous; partial, residual, post, comb, fn and the +//! outputs written in 16-byte stores (residualOut, layerInput) start on 16-byte boundaries. +struct MhcArBoundaryArgs +{ + void** workspace; //!< the TP group's TensorRT-LLM all-reduce workspace (Lamport buffers, flag, sizes) + int rank; //!< this rank in the group of 4 + void const* partial; //!< [M, 4096] bf16 this rank's partial output of the block + void const* residual; //!< [M, 4, 4096] bf16 the streams before the block's post-mapping + float const* postPrev; //!< [M, 4] + float const* combPrev; //!< [M, 4, 4] + float const* fn; //!< [24, 4 * 4096] the next block's mixer weights + float const* scale; //!< [3] + float const* base; //!< [24] + void const* normWeight; //!< [4096] bf16 the next block's input RMSNorm + void* residualOut; //!< [M, 4, 4096] bf16 + float* postOut; //!< [M, 4] + float* combOut; //!< [M, 4, 4] + void* layerInput; //!< [M, 4096] bf16 RMS-normalized + int rows; //!< M, 1 to kMhcArBoundaryMaxRows + int sinkhornIters; //!< at least 1 + float normEps; //!< the mixer's RMS epsilon + float hcEps; //!< added to pre + float sinkhornEps; //!< the Sinkhorn normalization's epsilon + float postMult; //!< post's multiplier + float layerInputEps; //!< the next block's RMSNorm epsilon + bool pdl; //!< programmatic dependent launch (only the weights are read before the wait) +}; + +//! One launch on `stream`: the all-reduce reads and rotates the workspace's Lamport state, so it is ordered with the +//! other users of that workspace (the library's one-shot all-reduce) by the stream. +//! +//! Returns cudaErrorInvalidValue for a row count out of range and otherwise the result of the CUDA calls. +cudaError_t mhcArBoundary(MhcArBoundaryArgs const& args, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_MV2_NAMESPACE_END diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/dims.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/dims.py new file mode 100644 index 000000000000..4f0576675f99 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/dims.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Every size the target computes with, read once from the text config and the mapping. + +The checkpoint is the composite ``Glm5NextForConditionalGeneration`` release; the text +model's fields sit on ``text_config``. Shares are this rank's under tensor parallelism 4: +attention heads (KDA and MLA query heads) and FFN intermediates are split over the four +ranks, the latent and indexer state and every norm are replicated. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from transformers import PretrainedConfig + +from tensorrt_llm.mapping import Mapping + +from .schedule import LayerSchedule + +#: The checkpoint's key prefix of the text model. +TEXT_PREFIX = "model.language_model." + + +@dataclass(frozen=True) +class Dims: + tp_size: int + tp_rank: int + hidden: int + vocab: int + num_layers: int # decoder layers; the MTP layer is index num_layers + schedule: LayerSchedule + rms_eps: float + # KDA (linear attention) + kda_heads: int # total + kda_head_dim: int + kda_conv_size: int + kda_gate_lower_bound: float + # Sparse MLA + mla_heads: int # total query heads + q_lora: int + kv_lora: int + qk_nope: int + v_head: int + index_heads: int + index_head_dim: int + index_kpool: int + index_topk: int + index_always_select_tail: bool + # FFN + intermediate: int # dense MLP, total + moe_intermediate: int # routed and shared expert, total + num_experts: int + top_k: int + routed_scaling: float + swiglu_limit: float + # mHC + hc_mult: int + hc_sinkhorn_iters: int + hc_eps: float + mtp: bool # the MTP layer is built (one-model MTP is configured) + max_num_tokens: int # the engine's per-step token budget (autotuning range of the MoE runner) + + # -- this rank's shares ------------------------------------------------------------ + + @property + def kda_heads_local(self) -> int: + return self.kda_heads // self.tp_size + + @property + def kda_channels_local(self) -> int: + return self.kda_heads_local * self.kda_head_dim + + @property + def mla_heads_local(self) -> int: + return self.mla_heads // self.tp_size + + @property + def intermediate_local(self) -> int: + return self.intermediate // self.tp_size + + @property + def moe_intermediate_local(self) -> int: + return self.moe_intermediate // self.tp_size + + @property + def vocab_local(self) -> int: + return -(-self.vocab // self.tp_size) + + @property + def select_k(self) -> int: + """Pools the indexer selects per row.""" + return self.index_topk // self.index_kpool + + @property + def mtp_layer(self) -> int: + return self.num_layers + + def layer_prefix(self, layer: int) -> str: + return f"{TEXT_PREFIX}layers.{layer}." + + @classmethod + def from_config( + cls, text: PretrainedConfig, mapping: Mapping, mtp: bool, max_num_tokens: int + ) -> "Dims": + """The sizes of a config ``routing.fp8_model_refused`` admits (the MoE routing, the + geometry the kernels compile in and the KDA gate bound are checked there).""" + linear = dict(text.linear_attn_config) + return cls( + tp_size=mapping.tp_size, + tp_rank=mapping.tp_rank, + hidden=int(text.hidden_size), + vocab=int(text.vocab_size), + num_layers=int(text.num_hidden_layers), + schedule=LayerSchedule.from_text_config(text), + rms_eps=float(text.rms_norm_eps), + kda_heads=int(linear["num_heads"]), + kda_head_dim=int(linear["head_dim"]), + kda_conv_size=int(linear["short_conv_kernel_size"]), + kda_gate_lower_bound=float(linear["gate_lower_bound"]), + mla_heads=int(text.num_attention_heads), + q_lora=int(text.q_lora_rank), + kv_lora=int(text.kv_lora_rank), + qk_nope=int(text.qk_nope_head_dim), + v_head=int(text.v_head_dim), + index_heads=int(text.index_n_heads), + index_head_dim=int(text.index_head_dim), + index_kpool=int(text.index_kpool), + index_topk=int(text.index_topk), + index_always_select_tail=bool(text.index_kpool_always_select_tail), + intermediate=int(text.intermediate_size), + moe_intermediate=int(text.moe_intermediate_size), + num_experts=int(text.n_routed_experts), + top_k=int(text.num_experts_per_tok), + routed_scaling=float(text.routed_scaling_factor), + swiglu_limit=float(text.swiglu_limit), + hc_mult=int(text.hc_mult), + hc_sinkhorn_iters=int(text.hc_sinkhorn_iters), + hc_eps=float(text.hc_eps), + mtp=mtp, + max_num_tokens=int(max_num_tokens), + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/fake_ops.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/fake_ops.py new file mode 100644 index 000000000000..b3d5ad1175c0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/fake_ops.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fake (meta) implementations of this target's native ops, built from its csrc/. + +The loader imports this module right after the target is found in the family +library, so a fake for an op the library does not define fails the load. Both ops +write their outputs in place and return nothing. Each fake is named ``fake_`` +and takes the op's schema arguments by their schema names, in schema order; the +target's fake-op test checks both against the library. +""" + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2._native import namespace_of + +_NS = namespace_of(__name__) + + +@torch.library.register_fake(f"{_NS}::decode_moe") +def fake_decode_moe( + x, logits, bias, scaling, limit, w31, s31, w2, s2, sw31, ss31, sw2, ss2, out, pdl +): + return None + + +@torch.library.register_fake(f"{_NS}::mhc_ar_boundary") +def fake_mhc_ar_boundary( + partial, + workspace, + rank, + residual, + post, + comb, + fn, + scale, + base, + norm_weight, + sinkhorn_iters, + norm_eps, + hc_eps, + sinkhorn_eps, + post_mult, + layer_input_eps, + residual_out, + post_out, + comb_out, + layer_input, + pdl, +): + return None diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/layers/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/layers/__init__.py new file mode 100644 index 000000000000..9c5daf39b19b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/layers/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The target's layers: KDA, sparse MLA, FFN and mHC, each on catalog and target operations.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/layers/ffn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/layers/ffn.py new file mode 100644 index 000000000000..cf08c578331b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/layers/ffn.py @@ -0,0 +1,277 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The FFN side: the three leading dense MLPs and the 288-expert MoE (top-8 plus one shared +expert), block-FP8 throughout, the intermediate split over the four ranks. + +Every projection is a 128 x 128 block-FP8 weight with FP32 block scales read as the +checkpoint stores them (``weight_scale_inv``, a dequantization multiplier), and every input +of it is the 1 x 128 block-FP8 quantization of its rows: ``fp8_quantize_1x128`` then +``fp8_blockwise_gemm``. The SwiGLU is clamped: ``min(gate, 10)`` and ``clamp(up, +-10)``. + +Layouts. A gated MLP (the dense layers and the shared expert) keeps its gate and up rows +stacked ``[gate; up]`` (the activation reads ``[gate | up]``). The routed experts are +``[E, 2I, H]`` with ``[up; gate]`` rows and ``[E, H, I]`` down weights, the plain layout the +TRTLLM-gen FP8 block-scale runner, ``decode_moe`` and ``prefill_moe`` all read; the decode +kernel also takes the shared expert in that layout, so the shared expert is kept in both. + +The MoE takes one of three paths by row count: + +* at most ``decode_moe.MAX_ROWS`` (4) rows: ``decode_moe``, routing, the routed experts and + the shared expert (as a ninth slot) in one persistent launch; +* more rows in a generation step: the router GEMM, the TRTLLM-gen runner (routing in the + kernel) and the shared expert's two GEMMs, summed; +* sequence-parallel prefill: ``prefill_moe`` on rows the caller quantized and gathered, with + the shared expert's output as its addend. + +Outputs are this rank's partials (each expert's K-split of its intermediate); the caller +reduces them once over the TP group. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from tensorrt_llm._torch.autotuner import AutoTuner + +from ..checkpoint import BLOCK, Checkpoint, cols, rows, share +from ..dims import Dims +from ..ops.activation.silu_and_mul import silu_and_mul +from ..ops.gemm.dsv3_router_gemm import dsv3_router_gemm +from ..ops.gemm.fp8_blockwise_gemm import fp8_blockwise_gemm +from ..ops.moe.decode_moe import MAX_ROWS as DECODE_MOE_ROWS +from ..ops.moe.decode_moe import decode_moe +from ..ops.moe.fp8_block_scale_moe_runner import ROUTING_DEEPSEEK_V3, fp8_block_scale_moe_runner +from ..ops.moe.prefill_moe import PrefillMoeWorkspace, prefill_moe +from ..ops.quantization.fp8_quantize_1x128 import fp8_quantize_1x128 + + +@dataclass +class GatedWeights: + """One TP-split gated MLP (a dense layer, or the shared expert).""" + + gate_up: torch.Tensor # [2 * I_local, H] e4m3, rows [gate; up] + gate_up_scale: torch.Tensor # [2 * I_local / 128, H / 128] fp32 + down: torch.Tensor # [H, I_local] e4m3 + down_scale: torch.Tensor # [H / 128, I_local / 128] fp32 + + +@dataclass +class DenseMlpWeights: + mlp: GatedWeights + + +@dataclass +class MoeWeights: + router: torch.Tensor # [E, H] bf16 + bias: torch.Tensor # [E] fp32, e_score_correction_bias + w31: torch.Tensor # [E, 2 * I_local, H] e4m3, rows [up; gate] + s31: torch.Tensor # [E, 2 * I_local / 128, H / 128] fp32 + w2: torch.Tensor # [E, H, I_local] e4m3 + s2: torch.Tensor # [E, H / 128, I_local / 128] fp32 + shared: GatedWeights + #: The shared expert in the routed experts' layout, decode_moe's ninth slot: + #: (w31 [2 * I_local, H] [up; gate], s31, w2 [H, I_local], s2). + shared_decode: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] + + +def _load_gated(ck: Checkpoint, prefix: str, intermediate: int, dims: Dims) -> GatedWeights: + """``prefix``'s gate / up / down projections at this rank's intermediate share: gate and up + rows, down columns, and the matching 128-blocks of their scales.""" + start, end = share(intermediate, dims.tp_size, dims.tp_rank) + sb, se = start // BLOCK, end // BLOCK + gate = ck.read(f"{prefix}gate_proj.weight", cut=rows(start, end)) + gate_s = ck.read(f"{prefix}gate_proj.weight_scale_inv", cut=rows(sb, se), dtype=torch.float32) + up = ck.read(f"{prefix}up_proj.weight", cut=rows(start, end)) + up_s = ck.read(f"{prefix}up_proj.weight_scale_inv", cut=rows(sb, se), dtype=torch.float32) + down = ck.read(f"{prefix}down_proj.weight", cut=cols(start, end)) + down_s = ck.read(f"{prefix}down_proj.weight_scale_inv", cut=cols(sb, se), dtype=torch.float32) + for t in (gate, up, down): + assert t.dtype == torch.float8_e4m3fn, ( + f"{prefix}: expected block-FP8 weights, got {t.dtype}" + ) + return GatedWeights( + gate_up=torch.cat([gate, up]).contiguous(), + gate_up_scale=torch.cat([gate_s, up_s]).contiguous(), + down=down, + down_scale=down_s, + ) + + +def load_dense_mlp(ck: Checkpoint, dims: Dims, layer: int) -> DenseMlpWeights: + return DenseMlpWeights( + _load_gated(ck, f"{dims.layer_prefix(layer)}mlp.", dims.intermediate, dims) + ) + + +def load_moe(ck: Checkpoint, dims: Dims, layer: int) -> MoeWeights: + prefix = f"{dims.layer_prefix(layer)}mlp." + num_experts, hidden = dims.num_experts, dims.hidden + inter = dims.moe_intermediate_local + start, end = share(dims.moe_intermediate, dims.tp_size, dims.tp_rank) + sb, se = start // BLOCK, end // BLOCK + device = ck.device + w31 = torch.empty((num_experts, 2 * inter, hidden), dtype=torch.float8_e4m3fn, device=device) + s31 = torch.empty( + (num_experts, 2 * inter // BLOCK, hidden // BLOCK), dtype=torch.float32, device=device + ) + w2 = torch.empty((num_experts, hidden, inter), dtype=torch.float8_e4m3fn, device=device) + s2 = torch.empty( + (num_experts, hidden // BLOCK, inter // BLOCK), dtype=torch.float32, device=device + ) + half = inter // BLOCK + for e in range(num_experts): + ep = f"{prefix}experts.{e}." + w31[e, :inter] = ck.read(f"{ep}up_proj.weight", cut=rows(start, end)) + w31[e, inter:] = ck.read(f"{ep}gate_proj.weight", cut=rows(start, end)) + s31[e, :half] = ck.read(f"{ep}up_proj.weight_scale_inv", cut=rows(sb, se)) + s31[e, half:] = ck.read(f"{ep}gate_proj.weight_scale_inv", cut=rows(sb, se)) + w2[e] = ck.read(f"{ep}down_proj.weight", cut=cols(start, end)) + s2[e] = ck.read(f"{ep}down_proj.weight_scale_inv", cut=cols(sb, se)) + shared = _load_gated(ck, f"{prefix}shared_experts.", dims.moe_intermediate, dims) + gate, up = shared.gate_up.chunk(2, dim=0) + gate_s, up_s = shared.gate_up_scale.chunk(2, dim=0) + return MoeWeights( + router=ck.read(f"{prefix}gate.weight", dtype=torch.bfloat16), + bias=ck.read(f"{prefix}gate.e_score_correction_bias", dtype=torch.float32), + w31=w31, + s31=s31, + w2=w2, + s2=s2, + shared=shared, + shared_decode=( + torch.cat([up, gate]).contiguous(), + torch.cat([up_s, gate_s]).contiguous(), + shared.down, + shared.down_scale, + ), + ) + + +def _gated_from_fp8( + w: GatedWeights, x_fp8: torch.Tensor, x_scale: torch.Tensor, limit: float +) -> torch.Tensor: + """A gated MLP's partial on rows already quantized 1 x 128 (``x_scale`` K-block-major): + gate_up, clamped SwiGLU, the activation's own 1 x 128 quantization, down.""" + h = fp8_blockwise_gemm(x_fp8, w.gate_up, x_scale, w.gate_up_scale) + act = silu_and_mul(h, swiglu_limit=limit) + a_fp8, a_scale = fp8_quantize_1x128(act) + return fp8_blockwise_gemm(a_fp8, w.down, a_scale, w.down_scale) + + +class DenseMlp: + def __init__(self, w: DenseMlpWeights, dims: Dims, layer_idx: int) -> None: + self.w = w.mlp + self.limit = dims.swiglu_limit + self.hidden = dims.hidden + self.layer_idx = layer_idx + + def partial(self, x: torch.Tensor) -> torch.Tensor: + flat = x.reshape(-1, self.hidden) + x_fp8, x_scale = fp8_quantize_1x128(flat) + return _gated_from_fp8(self.w, x_fp8, x_scale, self.limit).view_as(x) + + def partial_from_fp8(self, x_fp8: torch.Tensor, x_scale: torch.Tensor) -> torch.Tensor: + return _gated_from_fp8(self.w, x_fp8, x_scale, self.limit) + + +class MoeLayer: + def __init__(self, w: MoeWeights, dims: Dims, layer_idx: int) -> None: + self.w = w + self.hidden = dims.hidden + self.limit = dims.swiglu_limit + self.scaling = dims.routed_scaling + self.num_experts = dims.num_experts + self.top_k = dims.top_k + self.inter_local = dims.moe_intermediate_local + self.tune_max_num_tokens = dims.max_num_tokens + self.layer_idx = layer_idx + + def router_logits(self, x: torch.Tensor) -> torch.Tensor: + """FP32 ``[rows, E]`` router logits: the BF16 router GEMM with FP32 output.""" + return dsv3_router_gemm(x, self.w.router.t(), None, torch.float32) + + def partial(self, x: torch.Tensor) -> torch.Tensor: + """Routed experts plus the shared expert, this rank's partial, for a generation step.""" + w = self.w + flat = x.reshape(-1, self.hidden) + logits = self.router_logits(flat) + if flat.shape[0] <= DECODE_MOE_ROWS: + out = decode_moe( + flat.contiguous(), + logits, + w.bias, + self.scaling, + self.limit, + w.w31, + w.s31, + w.w2, + w.s2, + *w.shared_decode, + ) + return out.view_as(x) + x_fp8, x_scale = fp8_quantize_1x128(flat) + routed = self._runner(logits, x_fp8, x_scale) + shared = _gated_from_fp8(w.shared, x_fp8, x_scale, self.limit) + return (routed + shared).view_as(x) + + def _runner( + self, logits: torch.Tensor, x_fp8: torch.Tensor, x_scale: torch.Tensor + ) -> torch.Tensor: + """The routed experts of rows quantized 1 x 128 on the TRTLLM-gen runner (routing in + the kernel), this rank's partial.""" + w = self.w + return fp8_block_scale_moe_runner( + logits, + w.bias, + x_fp8, + x_scale, + w.w31, + w.s31, + w.w2, + w.s2, + self.num_experts, + self.top_k, + 0, # no shared expert fused into the routed weights + 1, # n_group + 1, # topk_group + self.inter_local, + 0, # local_expert_offset: every rank holds every expert + self.num_experts, + self.scaling, + ROUTING_DEEPSEEK_V3, + gemm1_clamp_limit=self.limit, + tune_max_num_tokens=self.tune_max_num_tokens, + ) + + def prefill_partial( + self, + x_fp8: torch.Tensor, + x_scale: torch.Tensor, + logits: torch.Tensor, + workspace: PrefillMoeWorkspace, + ) -> torch.Tensor: + """Every row of a sequence-parallel prefill step, from the gathered quantized rows and + router logits: the routed experts with the shared expert's output as the addend.""" + w = self.w + if AutoTuner.get().is_tuning_mode: + # Generation steps of more than decode_moe's rows run the TRTLLM-gen runner, which + # the engine's tuning warmup never reaches (its generation steps fit decode_moe). + # One call here tunes the runner for every row bucket up to the step token budget. + self._runner(logits, x_fp8, x_scale) + shared = _gated_from_fp8(w.shared, x_fp8, x_scale, self.limit) + return prefill_moe( + x_fp8, + x_scale, + logits, + w.bias, + self.scaling, + self.limit, + w.w31, + w.s31, + w.w2, + w.s2, + shared, + workspace, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/layers/kda.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/layers/kda.py new file mode 100644 index 000000000000..eb6ee6144684 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/layers/kda.py @@ -0,0 +1,412 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The KDA (Kimi delta attention) layers: this rank's heads, their recurrent state and conv +windows, and the row-split ``o_proj`` partial of every row of a step. + +Under tensor parallelism every tensor that feeds a head is split by this rank's heads (16 of 64 +at TP4, ``H`` heads of ``hd = 128`` channels, ``D = H * hd``): the ``q`` / ``k`` / ``v`` +projections and conv taps, ``b`` (beta), the up projections ``f_b`` / ``g_b`` of the low-rank +forget and output gates, ``A_log`` and ``dt_bias``. The gates' down projections ``f_a`` / +``g_a`` and the output norm are replicated, and ``o_proj`` keeps this rank's input columns, so +the layer's output is a partial the caller reduces. The checkpoint's KDA tensors are BF16 except +the FP32 ``A_log`` and ``dt_bias``, which stay FP32. + +Projections. One ``[q; k; v]`` GEMM gives the conv inputs. One ``[f_a; g_a; b]`` GEMM (rows +padded to a multiple of 8) gives both gates' low-rank inputs and the beta logits, and one batched +GEMM of its adjacent ``f_a`` / ``g_a`` columns with ``[f_b; g_b]`` gives the forget-gate and the +output-gate logits. ``[f_a; g_a; b]`` and ``o_proj`` run on ``low_m_bf16_gemm`` (its low-M +kernels for decode-sized steps, cuBLAS otherwise), the ``[q; k; v]`` GEMM on cuBLAS. In a CUDA-graph-captured +generation step of at most 128 rows the gate projections run on the aux stream beside the +``[q; k; v]`` GEMM. + +Rows of a step are packed context first, and each kind of row runs one recurrence kernel: + +* context rows: ``kda_prefill_conv`` (conv + SiLU from each slot's conv history, the history + written back), then ``kda_prefill`` from each slot's recurrent state (the persistent engine + route; slots of fresh requests are zeroed first by ``reset_recurrent_state_rows``), then the + gated output norm ``rms_norm_gated_token_major``; +* plain decode rows (one per request): ``kda_decode``, the conv, recurrence and gated output + norm in one launch, the conv window and state row updated in place; +* MTP verify rows (``1 + num_spec`` per request): ``kda_mtp_decode``, which replays the previous + round's accepted drafts, runs the new rows, commits the state and the conv window after the + golden (first) row and caches the new drafts for the next round. + +Cache contract (``kv_cache_manager.mamba_layer_cache(layer)``): ``conv`` is the live +``[slots, 3D, 3]`` BF16 conv pool (``[q | k | v]`` sections), ``temporal`` the +``[slots, H, 128, 128]`` V-first recurrent pool (FP32: the GLM cache manager allocates no other). With speculative +decoding the layer cache also carries the verify kernel's replay caches: their committed conv +window is seeded from the live pool (``commit_conv_window``) after every prefill and plain +decode, and ``prev_num_accepted_tokens`` (the previous round's accepted drafts per slot, which +the runtime records after sampling) tells the verify kernel how many cached drafts to replay. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional + +import torch + +from ..checkpoint import Checkpoint, Cut, cols, rows, share +from ..dims import Dims +from ..ops.gemm.low_m_bf16_gemm import low_m_bf16_gemm +from ..ops.norm.rms_norm_gated_token_major import rms_norm_gated_token_major +from ..ops.ssm.kda_decode import kda_decode +from ..ops.ssm.kda_mtp_decode import kda_mtp_decode +from ..ops.ssm.kda_prefill import kda_prefill +from ..ops.ssm.kda_prefill_conv import kda_prefill_conv +from ..ops.ssm.reset_recurrent_state_rows import reset_recurrent_state_rows +from ..schedule import StepSchedule +from .streams import AuxStream + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata + from tensorrt_llm._torch.attention.backends.sparse.glm_kpool.metadata import ( + Glm5NextMamba2Metadata, + ) + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaLayerCache + +#: The head size and conv width the KDA kernels implement. +_HEAD_DIM = 128 +_CONV_WIDTH = 4 +#: The ``[f_a; g_a; b]`` projection's rows are padded to a multiple of this. +_GATE_ROWS_ALIGN = 8 +#: Captured generation steps of at most this many rows run the gate projections on the aux +#: stream beside the ``[q; k; v]`` GEMM. +_OVERLAP_MAX_ROWS = 128 +#: KDA prefill takes the persistent engine route: the chunk-decomposed route, which the auto +#: router picks whenever 2 * sequences * heads fits the SMs (16 heads per rank and FP32 state +#: here), can stall one CTA of ``kernel_chain_dv2`` on its mbarrier pipeline under serving load +#: (alternating 1- and 8-request 13K-token prefills) and hang every rank at the next collective. +#: Return to the auto route once that route passes the same alternating serving load. +_PREFILL_ROUTE = "engine" + + +@dataclass +class KdaWeights: + """One KDA layer at this rank's heads (``H`` heads, ``D = H * 128`` channels), in the + layouts the operations read.""" + + qkv: torch.Tensor # [3D, hidden] bf16, rows [q; k; v] + fa_ga_b: torch.Tensor # [2 * 128 + H padded to 8, hidden] bf16, rows [f_a; g_a; b; zeros] + fb_gb: torch.Tensor # [2D, 128] bf16, rows [f_b; g_b] + #: ``fb_gb`` as the batched GEMM's ``[2, 128, D]`` operand ``(f_b^T, g_b^T)`` (a view). + fb_gb_t: torch.Tensor + o_proj: torch.Tensor # [hidden, D] bf16, this rank's input columns + a_log: torch.Tensor # [H] fp32 + dt_bias: torch.Tensor # [D] fp32 + o_norm: torch.Tensor # [128] bf16, the gated output norm's weight + o_norm_f32: torch.Tensor # [128] fp32, the same values for the fused decode / verify norms + conv: torch.Tensor # [3D, 4] bf16, the q, k and v taps stacked (prefill conv) + conv_t: tuple[torch.Tensor, torch.Tensor, torch.Tensor] # q, k, v taps [4, D] bf16 (decode) + conv_f32: tuple[torch.Tensor, torch.Tensor, torch.Tensor] # q, k, v taps [D, 4] fp32 (verify) + conv_bias: torch.Tensor # [D] bf16 zeros, the decode kernel's bias of the bias-free conv + + +def load_kda(ck: Checkpoint, dims: Dims, layer: int) -> KdaWeights: + """Read layer ``layer``'s KDA tensors at this rank's heads and build the operations' layouts. + + Every ``self_attn`` tensor of the layer is read once: rows of this rank's channels for + ``q / k / v_proj``, ``f_b / g_b_proj``, the conv taps and ``dt_bias``; rows of this rank's + heads for ``b_proj`` and ``A_log``; this rank's input columns of ``o_proj``; ``f_a / g_a_proj`` + and ``o_norm`` whole. Projections, taps and the norm are BF16, ``A_log`` and ``dt_bias`` FP32. + """ + assert dims.kda_head_dim == _HEAD_DIM and dims.kda_conv_size == _CONV_WIDTH, ( + f"the KDA kernels take head_dim {_HEAD_DIM} and conv width {_CONV_WIDTH}, not " + f"{dims.kda_head_dim} and {dims.kda_conv_size}" + ) + prefix = f"{dims.layer_prefix(layer)}self_attn." + hidden, hd = dims.hidden, dims.kda_head_dim + heads, d = dims.kda_heads_local, dims.kda_channels_local + h0, h1 = share(dims.kda_heads, dims.tp_size, dims.tp_rank) + channels = rows(h0 * hd, h1 * hd) + device = ck.device + bf16 = torch.bfloat16 + + def read(name: str, cut: Optional[Cut] = None, dtype: torch.dtype = bf16) -> torch.Tensor: + return ck.read(f"{prefix}{name}", cut=cut, dtype=dtype) + + qkv = torch.empty((3 * d, hidden), dtype=bf16, device=device) + conv = torch.empty((3 * d, _CONV_WIDTH), dtype=bf16, device=device) + for i, part in enumerate("qkv"): + qkv[i * d : (i + 1) * d] = read(f"{part}_proj.weight", channels) + taps = read(f"{part}_conv1d.weight", channels) + assert taps.shape == (d, 1, _CONV_WIDTH), f"{prefix}{part}_conv1d: {tuple(taps.shape)}" + conv[i * d : (i + 1) * d] = taps.squeeze(1) + num_gate_rows = 2 * hd + heads + fa_ga_b = torch.zeros( + (-(-num_gate_rows // _GATE_ROWS_ALIGN) * _GATE_ROWS_ALIGN, hidden), + dtype=bf16, + device=device, + ) + fa_ga_b[:hd] = read("f_a_proj.weight") + fa_ga_b[hd : 2 * hd] = read("g_a_proj.weight") + fa_ga_b[2 * hd : num_gate_rows] = read("b_proj.weight", rows(h0, h1)) + fb_gb = torch.empty((2 * d, hd), dtype=bf16, device=device) + fb_gb[:d] = read("f_b_proj.weight", channels) + fb_gb[d:] = read("g_b_proj.weight", channels) + o_norm = read("o_norm.weight") + assert o_norm.shape == (hd,), f"{prefix}o_norm: {tuple(o_norm.shape)}" + weights = KdaWeights( + qkv=qkv, + fa_ga_b=fa_ga_b, + fb_gb=fb_gb, + fb_gb_t=fb_gb.view(2, d, hd).transpose(1, 2), + o_proj=read("o_proj.weight", cols(h0 * hd, h1 * hd)), + # Own allocations: the verify kernel takes every operand as 16-byte aligned unchecked, and + # a read may be a view into the entry it was cut from. + a_log=read("A_log", rows(h0, h1), torch.float32).clone(), + dt_bias=read("dt_bias", channels, torch.float32).clone(), + o_norm=o_norm, + o_norm_f32=o_norm.float(), + conv=conv, + conv_t=tuple(t.t().contiguous() for t in conv.split(d)), + conv_f32=tuple(t.float().contiguous() for t in conv.split(d)), + conv_bias=torch.zeros(d, dtype=bf16, device=device), + ) + assert weights.o_proj.shape == (hidden, d), f"{prefix}o_proj: {tuple(weights.o_proj.shape)}" + assert weights.a_log.shape == (heads,) and weights.dt_bias.shape == (d,), ( + f"{prefix}A_log {tuple(weights.a_log.shape)}, dt_bias {tuple(weights.dt_bias.shape)}" + ) + return weights + + +def _commit_conv_window( + cache: MambaLayerCache, slots: torch.Tensor, conv_pool: torch.Tensor +) -> None: + """Seed the verify kernel's replay conv caches of ``slots`` from the live conv pool, which a + prefill or a plain decode just advanced (a layer cache without replay caches has none).""" + if getattr(cache, "has_kda_replay_caches", False): + cache.commit_conv_window(slots, conv_pool) + + +class KdaLayer: + """One KDA layer's forward: the unreduced ``o_proj`` partial of every row of a step. + + Args: + w: the layer's weights (``load_kda``). + dims: the model's sizes and this rank's shares. + layer_idx: the decoder layer's index, which names its recurrent and conv state in the + hybrid cache. + aux: the aux stream the gate projections of a captured generation step run on. + """ + + def __init__(self, w: KdaWeights, dims: Dims, layer_idx: int, aux: AuxStream) -> None: + self.w = w + self.layer_idx = layer_idx + self.heads = dims.kda_heads_local + self.head_dim = dims.kda_head_dim + self.channels = dims.kda_channels_local + self.lower_bound = dims.kda_gate_lower_bound + self.scale = dims.kda_head_dim**-0.5 + self.norm_eps = dims.rms_eps + self.aux = aux + + def partial( + self, x: torch.Tensor, step: StepSchedule, attn_metadata: AttentionMetadata + ) -> torch.Tensor: + """This rank's ``o_proj`` partial (before the TP reduction) of every row of the step. + + Args: + x: ``[N, hidden]`` bf16, the step's rows packed context first (``N`` is + ``step.num_ctx_tokens + step.num_generation_tokens``). + step: the step's schedule. + attn_metadata: the engine-prepared metadata: its ``mamba_metadata`` (slots, packed + context offsets, prefill schedules) and ``kv_cache_manager``. + + Returns: + ``[N, hidden]`` bf16. + """ + num_contexts, num_ctx_tokens = step.num_contexts, step.num_ctx_tokens + num_rows = num_ctx_tokens + step.num_generation_tokens + assert x.shape[0] == num_rows, f"{x.shape[0]} rows for a step of {num_rows}" + hybrid = attn_metadata.mamba_metadata + cache = attn_metadata.kv_cache_manager.mamba_layer_cache(self.layer_idx) + slots = hybrid.state_indices[: num_contexts + step.num_generations] + cores = [] + if num_contexts: + cores.append(self._prefill(x[:num_ctx_tokens], hybrid, cache, slots[:num_contexts])) + if step.num_generations: + # A 16-byte-aligned copy when the verify kernel needs one, else the slots' tail. + generation_slots = hybrid.generation_state_indices + if generation_slots is None: + generation_slots = slots[num_contexts:] + generation = x[num_ctx_tokens:num_rows] + tokens = step.gen_tokens_per_request + if tokens == 1: + cores.append(self._decode(generation, cache, generation_slots)) + else: + cores.append(self._verify(generation, tokens, cache, generation_slots)) + core = cores[0] if len(cores) == 1 else torch.cat(cores) + return low_m_bf16_gemm(core, self.w.o_proj) + + # -- projections ---------------------------------------------------------------------- + + def _heads(self, columns: torch.Tensor) -> torch.Tensor: + """``[n, D]`` columns as the kernels' ``[1, n, H, 128]`` view (the row stride kept).""" + return columns.unflatten(-1, (self.heads, self.head_dim)).unsqueeze(0) + + def _gates(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``(beta [n, H], forget [n, D], out_gate [n, D])`` logits of rows ``x [n, hidden]``: + the ``[f_a; g_a; b]`` GEMM, then ``f_b`` / ``g_b`` as one batched GEMM over its adjacent + ``f_a`` / ``g_a`` columns. ``beta`` is a column view of the first GEMM's output.""" + hd = self.head_dim + fa_ga_b = low_m_bf16_gemm(x, self.w.fa_ga_b) + beta = fa_ga_b[:, 2 * hd : 2 * hd + self.heads] + pair = fa_ga_b[:, : 2 * hd].reshape(-1, 2, hd).transpose(0, 1) # [2, n, hd] strided + gates = torch.bmm(pair, self.w.fb_gb_t) # [2, n, D] + return beta, gates[0], gates[1] + + def _project( + self, x: torch.Tensor + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + """``(qkv [n, 3D], (beta, forget, out_gate))`` of generation rows ``x``. In a captured + step of at most ``_OVERLAP_MAX_ROWS`` rows the gate projections run on the aux stream + while the ``[q; k; v]`` GEMM runs on the current one; otherwise both run in order.""" + qkv = lambda: torch.nn.functional.linear(x, self.w.qkv) # noqa: E731 + gates = lambda: self._gates(x) # noqa: E731 + if x.shape[0] > _OVERLAP_MAX_ROWS: + return qkv(), gates() + return self.aux.fork_join(qkv, gates) + + # -- the three row kinds ---------------------------------------------------------------- + + def _prefill( + self, + x: torch.Tensor, + hybrid: Glm5NextMamba2Metadata, + cache: MambaLayerCache, + slots: torch.Tensor, + ) -> torch.Tensor: + """The gated core ``[T, D]`` of the context rows ``x [T, hidden]`` of ``slots`` (int32 + ``[num_contexts]``), each sequence continuing from its slot's conv history and state.""" + w = self.w + tokens, num_contexts = x.shape[0], slots.shape[0] + conv_pool, state_pool = cache.conv, cache.temporal + has_init = hybrid.has_initial_states[:num_contexts] + qkv = torch.nn.functional.linear(x, w.qkv) + # The recurrence reads every slot as its initial state, so fresh slots start from zero; + # the prefill conv clears fresh conv histories itself. + reset_recurrent_state_rows( + recurrent_states=state_pool, state_indices=slots, has_initial_states=has_init + ) + q, k, v = kda_prefill_conv( + src=qkv, + weight=w.conv, + conv_state=conv_pool, + starts=hybrid.query_start_loc[: num_contexts + 1], + slots=slots, + has_init=has_init, + max_seqlen=hybrid.kda_max_seqlen, + ) + beta, forget, out_gate = self._gates(x) + o = kda_prefill( + q=q, + k=k, + v=v, + g=self._heads(forget), + beta=beta.unsqueeze(0), + state_pool=state_pool, + state_indices=slots, + cu_seqlens=hybrid.query_start_loc_long[: num_contexts + 1], + A_log=w.a_log, + dt_bias=w.dt_bias, + scale=self.scale, + lower_bound=self.lower_bound, + mode=_PREFILL_ROUTE, + max_seqlen=hybrid.kda_max_seqlen, + seq_order=hybrid.kda_seq_order, + cu_chunks=hybrid.kda_cu_chunks, + sm_reserve=0, # feeds the Small-BH route only + ) + _commit_conv_window(cache, slots, conv_pool) + core = rms_norm_gated_token_major( + o.view(-1, self.head_dim), + out_gate.view(tokens, self.heads, self.head_dim), + w.o_norm, + self.norm_eps, + gate_activation="sigmoid", + ) + return core.view(tokens, self.channels) + + def _decode(self, x: torch.Tensor, cache: MambaLayerCache, slots: torch.Tensor) -> torch.Tensor: + """The gated core ``[B, D]`` of one new row per request (``x [B, hidden]``, ``slots`` + int32 ``[B]``); the requests' conv windows and FP32 state rows advance in place.""" + w = self.w + d, batch = self.channels, x.shape[0] + conv_pool, state_pool = cache.conv, cache.temporal + assert state_pool.dtype == torch.float32, f"KDA state pool {state_pool.dtype}, not FP32" + qkv, (beta, forget, out_gate) = self._project(x) + # Allocated per call: inside a CUDA graph capture this is the graph pool's stable buffer. + out = x.new_empty((batch, 1, self.heads, self.head_dim)) + kda_decode( + x_q=self._heads(qkv[:, :d]), + x_k=self._heads(qkv[:, d : 2 * d]), + x_v=self._heads(qkv[:, 2 * d : 3 * d]), + w_q_t=w.conv_t[0], + w_k_t=w.conv_t[1], + w_v_t=w.conv_t[2], + bias_q=w.conv_bias, + bias_k=w.conv_bias, + bias_v=w.conv_bias, + conv_state_q=conv_pool[:, :d], + conv_state_k=conv_pool[:, d : 2 * d], + conv_state_v=conv_pool[:, 2 * d :], + a_log=w.a_log, + g=self._heads(forget), + dt_bias=w.dt_bias, + beta=beta.unsqueeze(0), + onorm_g=self._heads(out_gate), + onorm_weight=w.o_norm_f32, + ssm_state_indices=slots, + state=state_pool, + lower_bound=self.lower_bound, + scale=self.scale, + onorm_eps=self.norm_eps, + output=out, + ) + _commit_conv_window(cache, slots, conv_pool) + return out.view(batch, d) + + def _verify( + self, x: torch.Tensor, tokens_per_request: int, cache: MambaLayerCache, slots: torch.Tensor + ) -> torch.Tensor: + """The gated core ``[B * T, D]`` of ``T = tokens_per_request`` rows per request (the + golden row, then ``T - 1`` drafts; ``x [B * T, hidden]``, ``slots`` int32 ``[B]``).""" + if not getattr(cache, "has_kda_replay_caches", False): + raise RuntimeError( + f"KDA layer {self.layer_idx}: a verify step needs the KDA replay caches, which the " + "cache manager allocates for speculative decoding the fused verify kernel supports" + ) + w = self.w + d = self.channels + qkv, (beta, forget, out_gate) = self._project(x) + out = kda_mtp_decode( + x_q=self._heads(qkv[:, :d]), + x_k=self._heads(qkv[:, d : 2 * d]), + x_v=self._heads(qkv[:, 2 * d : 3 * d]), + w_q=w.conv_f32[0], + w_k=w.conv_f32[1], + w_v=w.conv_f32[2], + cs_q=cache.kda_conv_q, + cs_k=cache.kda_conv_k, + cs_v=cache.kda_conv_v, + g=self._heads(forget), + beta=beta.unsqueeze(0), + A_log=w.a_log, + dt_bias=w.dt_bias, + recurrent_state=cache.temporal, + k_cache=cache.kda_k_cache, + g_cache=cache.kda_g_cache, + v_cache=cache.kda_v_cache, + beta_cache=cache.kda_beta_cache, + ssm_state_indices=slots, + num_accepted_tokens=cache.prev_num_accepted_tokens, + num_spec=tokens_per_request - 1, + lower_bound=self.lower_bound, + scale=self.scale, + onorm_g=self._heads(out_gate), + onorm_weight=w.o_norm_f32, + onorm_eps=self.norm_eps, + ) + return out.view(x.shape[0], d) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/layers/mhc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/layers/mhc.py new file mode 100644 index 000000000000..8f8bc6331a82 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/layers/mhc.py @@ -0,0 +1,122 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One manifold hyper-connection (mHC) site of GLM-5.3-Flash: 4 residual streams around each +attention and FFN sublayer. + +Before a sublayer the streams are mixed (``pre_mapping``): a ``[24, 4 * hidden]`` mixer over +the RMS-normalized flattened streams gives the sublayer's input weights ``pre``, the output +weights ``post`` and a Sinkhorn-normalized 4 x 4 stream combination ``comb``. After it the +sublayer's output is written back (``post_mapping``: ``streams' = comb^T streams + post * +out``). ``fused_hc`` is a sublayer's post-mapping fused with the next one's pre-mapping and the +next sublayer's input RMSNorm. + +The parameters are the model's (``glm5_next_hyper_connection`` in the checkpoint's release): +the mixer's RMS epsilon is ``rms_norm_eps``, ``pre``'s epsilon and the Sinkhorn epsilon are +``hc_eps``, ``post``'s multiplier is 2 and the Sinkhorn runs ``hc_sinkhorn_iters`` rounds. +``fused_hc`` takes two epsilons: ``Mhc.mixer_eps`` is the mixer's, and its ``norm_eps`` +argument is the next sublayer's input RMSNorm epsilon (``layer_input_eps`` of the ops). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from ..dims import Dims +from ..ops.mhc.mhc_big_fuse import big_fuse_block_size, mhc_big_fuse +from ..ops.mhc.mhc_fused_hc import fused_hc_scratch, fused_hc_tactic, mhc_fused_hc +from ..ops.mhc.mhc_gemm_sqrsum import gemm_sqrsum_tile_n, mhc_gemm_sqrsum +from ..ops.mhc.mhc_post_mapping import mhc_post_mapping + +if TYPE_CHECKING: + from ..weights import HcWeights + +#: post's multiplier: post = sigmoid(.) * 2. +_POST_MULT = 2.0 + + +class Mhc: + """One mHC site's weights and the three operations on them. + + ``post`` travels as ``[rows, mult]`` FP32 and ``comb`` as ``[rows, mult, mult]`` FP32; the + streams are ``[rows, mult, hidden]`` BF16. + """ + + def __init__(self, w: HcWeights, dims: Dims) -> None: + self.fn = w.fn # [24, mult * hidden] fp32 + self.scale = w.scale # [3] fp32 + self.base = w.base # [24] fp32 + self.mult = dims.hc_mult + self.sinkhorn_iters = dims.hc_sinkhorn_iters + self.mixer_eps = dims.rms_eps # the mixer's RMS epsilon + self.eps = dims.hc_eps + self.sinkhorn_eps = dims.hc_eps + self.post_mult_value = _POST_MULT + + def scalars(self, layer_input_eps: float) -> tuple[int, float, float, float, float, float]: + """The scalar arguments the mHC operations take after the weights, in their order: + the Sinkhorn rounds, the mixer's RMS epsilon, ``pre``'s epsilon, the Sinkhorn epsilon, + ``post``'s multiplier and the next sublayer's input RMSNorm epsilon.""" + return ( + self.sinkhorn_iters, + self.mixer_eps, + self.eps, + self.sinkhorn_eps, + self.post_mult_value, + layer_input_eps, + ) + + def pre_mapping( + self, residual: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``(post, comb, x)``: the first sublayer's stream mix, ``x`` not yet RMS-normalized.""" + rows = residual.shape[0] + mixes, sqrsum = mhc_gemm_sqrsum(residual.view(rows, -1), self.fn, gemm_sqrsum_tile_n(rows)) + return mhc_big_fuse( + mixes, + sqrsum, + residual, + self.scale, + self.base, + None, + *self.scalars(0.0), + big_fuse_block_size(rows), + ) + + def fused_hc( + self, + out: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + norm_weight: torch.Tensor, + norm_eps: float, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """``(residual, post, comb, x)``: ``out``'s post-mapping into the streams, then the next + sublayer's pre-mapping and its input RMSNorm (``norm_weight``, ``norm_eps``).""" + args = ( + out, + residual, + post, + comb, + self.fn, + self.scale, + self.base, + norm_weight, + *self.scalars(norm_eps), + ) + tactic = fused_hc_tactic(*args) + return mhc_fused_hc(*args, tactic, *fused_hc_scratch(out.shape[0], tactic, out.device)) + + def post_mapping( + self, + out: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + norm_weight: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """The streams after ``out``'s post-mapping (the last sublayer's).""" + del norm_weight + return mhc_post_mapping(out, residual, post, comb) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/layers/mla.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/layers/mla.py new file mode 100644 index 000000000000..143b10a61848 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/layers/mla.py @@ -0,0 +1,701 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The sparse MLA layers: NoPE latent attention over a pool-compressed sparse selection. + +A row at cache position ``p`` attends its own prefix sparsely: the ``select_k`` (512) pools of +``index_kpool`` (4) consecutive positions an indexer scores highest among the row's complete +pools, plus the positions of its incomplete trailing pool. The cache holds per position the +RMS-normalized ``kv_a`` latent (512 wide, FP8, no rotary part) and the indexer row ``[k | gate | +pool key]`` (BF16; the pool key lives in a pool's first-member row). Queries and values are +absorbed into the latent space: head ``h`` scores latent ``c`` as ``(q_h @ w_k[h]) . c`` and maps +its latent output ``o`` to ``o @ w_v_t[h]``, ``w_k`` / ``w_v_t`` being the head's key and value +rows of ``kv_b_proj``; the softmax scale is ``qk_nope ** -0.5``. + +Weights (``load_mla``). The ``[q_a; kv_a]`` latent projections (one fused 128 x 128 block-FP8 +GEMM), the norms and the whole indexer (its ``[wk; compress gate; weights_proj]`` rows fused into +one BF16 GEMM, and ``wq_b``) are replicated. ``q_b_proj``, the absorbed ``kv_b_proj`` and +``o_proj``'s input columns exist twice: this rank's 16 of the 64 query heads (generation rows, +head parallel) and all 64 (sequence-parallel context rows). Products of at most 32 rows with +``[q_a; kv_a]`` or this rank's ``q_b`` / ``o_proj`` run on ``low_m_bf16_gemm`` over a dequantized +BF16 copy of the weight, without quantizing the activation; larger ones, and the all-head +products, quantize the rows 1 x 128 and run the blockwise FP8 GEMM. + +Generation rows (``generation_partial``; decode is ``T = 1``, an MTP verify ``T = 4``): the two +input GEMMs, then ``mla_prologue`` (the norms and both cache writes in one launch), then two +branches: the selection (``kpool_update`` of the rows' pools, the indexer query, ``kpool_score`` +and ``kpool_select``: top-k, expansion to latent rows and packing in one launch) and the query +(``q_b``, the absorbed key product, ``saturate_fp8``). In a CUDA graph capture the second input +GEMM and the query branch run on the aux stream. Then ``sparse_mla_decode`` with this rank's +heads, ``latent_values`` (zero for a row that selected nothing) and ``o_proj``'s row-parallel +partial, which the caller reduces. + +Context rows (``context_rows``): every rank appends every context row's state +(``mla_context_prologue``) and refreshes the pools the step completes (``kpool_update`` at the +pool-final positions), then attends only its own rows ``[lo, hi)`` with all 64 heads: the FP8 pool +scores (``kpool_score_fp8``, in calls of bounded memory), ``indexer_topk``, ``kpool_expand``, +``compact_sparse_rows``, the all-head ``q_b`` and key product, ``sparse_mla_decode``, the value +product quantized 1 x 128 (``quantize_1x128``) and the all-head ``o_proj``: the rows' full output. + +Cache contract: ``kv_cache_manager.get_latent_state_buffer(layer)`` (``[slots, tokens_per_block, +1, kv_lora]`` FP8) and ``get_index_state_buffer(layer)`` (``[slots, tokens_per_block, 1, 3 * +index_head_dim]`` BF16) are slot-major views of coalesced paged storage, addressed through the +step's GLM block tables. The sparse attention reads the latent storage as uniformly strided rows: +cache position ``p`` of table ``t`` is row ``base_row + tables[t, p // tokens_per_block] * +rows_per_slot + p % tokens_per_block``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional + +import torch + +from ..checkpoint import BLOCK, Checkpoint, dequantize_block_fp8, share +from ..dims import Dims +from ..ops.attention.compact_sparse_rows import compact_sparse_rows +from ..ops.attention.head_major import quantize_1x128, saturate_fp8 +from ..ops.attention.indexer_topk import indexer_topk +from ..ops.attention.kpool_expand import kpool_expand +from ..ops.attention.kpool_score import kpool_score +from ..ops.attention.kpool_score_fp8 import kpool_score_fp8 +from ..ops.attention.kpool_select import kpool_select +from ..ops.attention.kpool_update import kpool_update +from ..ops.attention.latent_values import latent_values +from ..ops.attention.mla_context_prologue import mla_context_prologue +from ..ops.attention.mla_prologue import mla_prologue +from ..ops.attention.sparse_mla_decode import ( + KV_LORA_RANK, + PAGE_ROWS, + WORKSPACE_BYTES, + counter_bytes, + sparse_mla_decode, +) +from ..ops.gemm.fp8_blockwise_gemm import fp8_blockwise_gemm +from ..ops.gemm.low_m_bf16_gemm import low_m_bf16_gemm +from ..ops.quantization.fp8_quantize_1x128 import fp8_quantize_1x128 +from ..schedule import StepSchedule +from .streams import AuxStream + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention.backends.sparse.glm_kpool.cache_manager import ( + Glm5NextCacheManager, + ) + +#: Block-FP8 projections of at most this many rows run on ``low_m_bf16_gemm`` over the +#: weight's dequantized BF16 copy; more rows run the FP8 GEMM. +_BF16_COPY_MAX_ROWS = 32 +#: The indexer key LayerNorm's epsilon. +_INDEX_KEY_NORM_EPS = 1e-6 +#: The expanded selection is padded to the sparse attention's 64-column tiles. +_SELECTION_ALIGN = 64 +#: FP32 elements one context scoring call may hold (the padded logits, a several-table call's +#: row copy and the gathered pool keys); larger calls are split. +_SCORE_ELEMENTS = 1 << 26 +#: The FP8 latent cache's dequantization scale: the cache is stored unscaled. +_KV_SCALE = 1.0 + + +@dataclass +class MlaHeads: + """The per-head projections of a set of ``H`` query heads: this rank's 16 (generation rows) + or all 64 (sequence-parallel context rows).""" + + q_b: torch.Tensor # [H * qk_nope, q_lora] e4m3 + q_b_scale: torch.Tensor # [H * qk_nope / 128, q_lora / 128] fp32 + w_k: torch.Tensor # [H, qk_nope, kv_lora] bf16: kv_b_proj's key rows of each head + w_v_t: torch.Tensor # [H, kv_lora, v_head] bf16: its value rows, transposed + o: torch.Tensor # [hidden, H * v_head] e4m3: o_proj's input columns of the heads + o_scale: torch.Tensor # [hidden / 128, H * v_head / 128] fp32 + + +@dataclass +class MlaWeights: + """One sparse MLA layer at this rank's share. Block-FP8 weights keep the checkpoint's + 128 x 128 FP32 block scales (``weight_scale_inv``) as stored.""" + + qkv_a: torch.Tensor # [q_lora + kv_lora, hidden] e4m3, rows [q_a; kv_a] + qkv_a_scale: torch.Tensor # [(q_lora + kv_lora) / 128, hidden / 128] fp32 + qkv_a_bf16: torch.Tensor # [q_lora + kv_lora, hidden] bf16, qkv_a dequantized + q_norm: torch.Tensor # [q_lora] bf16, q_a_layernorm + kv_norm: torch.Tensor # [kv_lora] bf16, kv_a_layernorm + #: [2 * index_head_dim + index_heads, hidden] bf16, rows [wk; compress gate; weights_proj] + index_kgw: torch.Tensor + index_q: torch.Tensor # [index_heads * index_head_dim, q_lora] bf16, wq_b + index_k_norm: torch.Tensor # [index_head_dim] bf16, the key LayerNorm's weight + index_k_bias: torch.Tensor # [index_head_dim] bf16, the key LayerNorm's bias + index_ape: torch.Tensor # [index_kpool, index_head_dim] bf16, the pool compression bias + #: This rank's heads; ``q_b``, ``q_b_scale``, ``w_k`` and ``w_v_t`` are views of ``full``'s. + local: MlaHeads + full: MlaHeads + q_b_bf16: torch.Tensor # [heads_local * qk_nope, q_lora] bf16, local.q_b dequantized + o_bf16: torch.Tensor # [hidden, heads_local * v_head] bf16, local.o dequantized + + +def load_mla(ck: Checkpoint, dims: Dims, layer: int) -> MlaWeights: + """Read layer ``layer``'s ``self_attn.*`` tensors: a decoder MLA layer or the MTP layer. + + Every tensor is read once and whole. The latent and indexer projections and the norms are + replicated; ``q_b_proj``, ``kv_b_proj`` and ``o_proj`` give the all-head copies, and this + rank's head shard is cut from them (``q_b`` rows, ``kv_b`` rows, ``o_proj`` columns). + """ + prefix = f"{dims.layer_prefix(layer)}self_attn." + heads, nope, v_head = dims.mla_heads, dims.qk_nope, dims.v_head + q_lora, kv_lora, hidden = dims.q_lora, dims.kv_lora, dims.hidden + index_heads, index_dim = dims.index_heads, dims.index_head_dim + + def read_bf16(name: str, shape: tuple[int, ...]) -> torch.Tensor: + tensor = ck.read(f"{prefix}{name}") + assert tensor.dtype == torch.bfloat16 and tuple(tensor.shape) == shape, ( + f"{prefix}{name}: {tensor.dtype} {tuple(tensor.shape)}, expected bf16 {shape}" + ) + return tensor + + def read_fp8(name: str, shape: tuple[int, int]) -> tuple[torch.Tensor, torch.Tensor]: + weight = ck.read(f"{prefix}{name}.weight") + scale = ck.read(f"{prefix}{name}.weight_scale_inv", dtype=torch.float32) + blocks = (-(-shape[0] // BLOCK), -(-shape[1] // BLOCK)) + assert weight.dtype == torch.float8_e4m3fn and tuple(weight.shape) == shape, ( + f"{prefix}{name}: {weight.dtype} {tuple(weight.shape)}, expected block-FP8 {shape}" + ) + assert tuple(scale.shape) == blocks, ( + f"{prefix}{name}: block scales {tuple(scale.shape)}, expected {blocks}" + ) + return weight, scale + + q_a, q_a_scale = read_fp8("q_a_proj", (q_lora, hidden)) + kv_a, kv_a_scale = read_fp8("kv_a_proj_with_mqa", (kv_lora, hidden)) + q_b, q_b_scale = read_fp8("q_b_proj", (heads * nope, q_lora)) + o, o_scale = read_fp8("o_proj", (hidden, heads * v_head)) + kv_b = read_bf16("kv_b_proj.weight", (heads * (nope + v_head), kv_lora)) + # The fused [q_a; kv_a] rows keep each source's 128-row scale blocks. + assert q_lora % BLOCK == 0, f"q_lora {q_lora} splits a 128-row scale block" + qkv_a = torch.cat([q_a, kv_a]) + qkv_a_scale = torch.cat([q_a_scale, kv_a_scale]) + + per_head = kv_b.view(heads, nope + v_head, kv_lora) + full = MlaHeads( + q_b=q_b, + q_b_scale=q_b_scale, + w_k=per_head[:, :nope].contiguous(), + w_v_t=per_head[:, nope:].transpose(1, 2).contiguous(), + o=o, + o_scale=o_scale, + ) + del per_head, kv_b + h0, h1 = share(heads, dims.tp_size, dims.tp_rank) + q0, q1, v0, v1 = h0 * nope, h1 * nope, h0 * v_head, h1 * v_head + assert q0 % BLOCK == 0 and v0 % BLOCK == 0, "a head share splits a 128-wide scale block" + local = MlaHeads( + q_b=q_b[q0:q1], + q_b_scale=q_b_scale[q0 // BLOCK : q1 // BLOCK], + w_k=full.w_k[h0:h1], + w_v_t=full.w_v_t[h0:h1], + o=o[:, v0:v1].contiguous(), + o_scale=o_scale[:, v0 // BLOCK : v1 // BLOCK].contiguous(), + ) + + index_kgw = torch.cat( + [ + read_bf16("indexer.wk.weight", (index_dim, hidden)), + read_bf16("indexer.index_kpool_compress_gate", (index_dim, hidden)), + read_bf16("indexer.weights_proj.weight", (index_heads, hidden)), + ] + ) + return MlaWeights( + qkv_a=qkv_a, + qkv_a_scale=qkv_a_scale, + qkv_a_bf16=dequantize_block_fp8(qkv_a, qkv_a_scale), + q_norm=read_bf16("q_a_layernorm.weight", (q_lora,)), + kv_norm=read_bf16("kv_a_layernorm.weight", (kv_lora,)), + index_kgw=index_kgw, + index_q=read_bf16("indexer.wq_b.weight", (index_heads * index_dim, q_lora)), + index_k_norm=read_bf16("indexer.k_norm.weight", (index_dim,)), + index_k_bias=read_bf16("indexer.k_norm.bias", (index_dim,)), + index_ape=read_bf16("indexer.index_kpool_compress_ape", (dims.index_kpool, index_dim)), + local=local, + full=full, + q_b_bf16=dequantize_block_fp8(local.q_b, local.q_b_scale), + o_bf16=dequantize_block_fp8(local.o, local.o_scale), + ) + + +class SparseAttentionScratch: + """Buffers every sparse MLA layer of the model shares, allocated before any CUDA graph capture. + + Every sparse attention launch of the model runs on the calling stream, one after another, so + one workspace and one set of semaphores serve them all: + + - ``workspace``: the sparse attention kernel's split-KV workspace, ``WORKSPACE_BYTES``, zeroed + once; + - ``counter(rows, heads)``: its multi-CTA semaphores, zeroed once (the kernel resets them after + every launch), sized for the engine's token budget: ``max_num_tokens`` generation rows of this + rank's heads, or a quarter of them as sequence-parallel context rows of all heads; + - ``kv_scale`` / ``kv_scale_inv``: the FP8 latent cache's dequantization and quantization scales + (FP32 ``[1]``), and ``kv_scale_value``, the former as a host float. + + Args: + dims: the model's sizes; ``max_num_tokens`` and the head counts size the semaphores. + """ + + def __init__(self, dims: Dims) -> None: + self.device = torch.device("cuda", torch.cuda.current_device()) + self.sm_count = torch.cuda.get_device_properties(self.device).multi_processor_count + self.workspace = torch.zeros(WORKSPACE_BYTES, dtype=torch.uint8, device=self.device) + # (row, head) pairs of the largest launch: every row of a step with this rank's heads, or + # a sequence-parallel quarter of the rows with all heads. + rows = dims.max_num_tokens + pairs = max(rows * dims.mla_heads_local, -(-rows // dims.tp_size) * dims.mla_heads) + self._counter = torch.zeros( + counter_bytes(pairs, 1, self.sm_count), dtype=torch.uint8, device=self.device + ) + self._outgrown: list[torch.Tensor] = [] + self.kv_scale_value = _KV_SCALE + self.kv_scale = torch.full((1,), _KV_SCALE, dtype=torch.float32, device=self.device) + self.kv_scale_inv = torch.full( + (1,), 1.0 / _KV_SCALE, dtype=torch.float32, device=self.device + ) + + def counter(self, rows: int, heads: int) -> torch.Tensor: + """The semaphores of a launch of ``rows`` query rows of ``heads`` heads (uint8, zero + between launches). A larger launch outside a capture gets a larger buffer; the outgrown + one stays allocated, since graphs captured on it keep addressing it.""" + need = counter_bytes(rows, heads, self.sm_count) + if self._counter.numel() < need: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"sparse MLA semaphores for {rows} rows x {heads} heads ({need} bytes) exceed " + f"the {self._counter.numel()} allocated before CUDA graph capture" + ) + self._outgrown.append(self._counter) + self._counter = torch.zeros(need, dtype=torch.uint8, device=self.device) + return self._counter + + +@dataclass(frozen=True) +class _CachePools: + """One sparse layer's paged cache as the kernels address it. + + ``latent`` and ``index`` are the slot-major pools; ``kv_rows`` views the latent storage as + uniformly strided rows (a whole number of 32-row pages), where cache position ``p`` of a + request is row ``base_row + slot(p) * rows_per_slot + p % tokens_per_block``. + """ + + latent: torch.Tensor # [slots, tokens_per_block, kv_lora] e4m3 + index: torch.Tensor # [slots, tokens_per_block, 3 * index_head_dim] bf16 + kv_rows: torch.Tensor # [N, 1, kv_lora] e4m3 + base_row: int + rows_per_slot: int + tokens_per_block: int + + +def _cache_pools(manager: Glm5NextCacheManager, layer_idx: int) -> _CachePools: + """Layer ``layer_idx``'s pools from the KV cache manager's documented buffers: views only, no + allocation or device work, so a captured step reads the same storage on replay.""" + latent = manager.get_latent_state_buffer(layer_idx) + index = manager.get_index_state_buffer(layer_idx) + if latent is None or index is None: + raise ValueError( + f"layer {layer_idx} has no latent / indexer pool in the KV cache manager; the layer " + "schedule and the cache layout disagree" + ) + latent, index = latent[:, :, 0, :], index[:, :, 0, :] + if latent.dtype != torch.float8_e4m3fn: + raise ValueError( + "the sparse MLA attends an FP8 latent cache (kv_cache_config.dtype fp8), got " + f"{latent.dtype}" + ) + if index.dtype != torch.bfloat16: + raise ValueError(f"the indexer cache is BF16, got {index.dtype}") + tokens_per_block = int(manager.tokens_per_block) + dim = latent.shape[-1] + if ( + latent.shape[1] != tokens_per_block + or index.shape[1] != tokens_per_block + or latent.stride(2) != 1 + or latent.stride(1) != dim + ): + raise ValueError( + f"layer {layer_idx}: latent pool {tuple(latent.shape)} / {latent.stride()} is not " + f"pages of {tokens_per_block} contiguous rows" + ) + slot_stride, offset = latent.stride(0), latent.storage_offset() + if slot_stride % dim or offset % dim: + raise ValueError( + f"layer {layer_idx}: latent slot stride {slot_stride} / storage offset {offset} are " + f"not multiples of the row width {dim}; the storage has no uniform row view" + ) + base_row, rows_per_slot = offset // dim, slot_stride // dim + total = latent.untyped_storage().nbytes() // (dim * latent.element_size()) + total -= total % PAGE_ROWS # the attention reads the rows as 32-row pages + top = base_row + (latent.shape[0] - 1) * rows_per_slot + tokens_per_block + if top > total: + raise ValueError( + f"layer {layer_idx}: latent rows up to {top} do not fit the {total} whole 32-row " + "pages of the cache storage" + ) + kv_rows = torch.as_strided(latent, (total, 1, dim), (dim, dim, 1), storage_offset=0) + return _CachePools(latent, index, kv_rows, base_row, rows_per_slot, tokens_per_block) + + +def _score_cost(rows: int, tables: int, pools: int, head_dim: int) -> int: + """FP32 elements ``kpool_score_fp8`` holds at once for ``rows`` rows over ``tables`` tables + of ``pools`` pools: the padded logits, a several-table call's row copy, and per gathered key + its BF16 row, FP8 row, scale and slot index.""" + if pools == 0: + return 0 + logits = -(-tables * pools // 256) * 256 + 256 + copy = max(1024, -(-pools // 64) * 64) + pools if tables > 1 else 0 + keys = tables * pools * (head_dim // 2 + head_dim // 4 + 3) + return rows * (logits + copy) + keys + 6 * pools + + +def _score_calls( + step: StepSchedule, lo: int, hi: int, kpool: int, head_dim: int +) -> list[tuple[int, int, int, int, int]]: + """The FP8 scoring calls of the packed context rows ``[lo, hi)``, each within + ``_SCORE_ELEMENTS``: ``(row_start, row_end, table_start, table_end, max_kv_len)``. + + Consecutive requests share a call while it fits; a request that does not fit alone is split + by rows. ``max_kv_len`` is the largest visible length among the call's rows (the position + 1 + of each request's last row in it). Host ints only. + """ + cu, cached = step.ctx_cu_seqlens, step.cached_lens + calls: list[tuple[int, int, int, int, int]] = [] + group: Optional[list[int]] = None # [row_start, row_end, table_start, table_end, max_kv_len] + + def cost(rows: int, tables: int, visible: int) -> int: + return _score_cost(rows, tables, visible // kpool, head_dim) + + for table in range(step.num_contexts): + first_row = int(cu[table]) + start, end = max(first_row, lo), min(int(cu[table + 1]), hi) + if start >= end: + continue + visible = int(cached[table]) + end - first_row + if group is not None: + merged = max(group[4], visible) + rows = group[1] - group[0] + end - start + if cost(rows, table + 1 - group[2], merged) <= _SCORE_ELEMENTS: + group[1], group[3], group[4] = end, table + 1, merged + continue + calls.append((group[0], group[1], group[2], group[3], group[4])) + group = None + if cost(end - start, 1, visible) <= _SCORE_ELEMENTS: + group = [start, end, table, table + 1, visible] + continue + fixed = cost(0, 1, visible) + tile = max(1, (_SCORE_ELEMENTS - fixed) // max(1, cost(1, 1, visible) - fixed)) + for row in range(start, end, tile): + row_end = min(row + tile, end) + calls.append((row, row_end, table, table + 1, int(cached[table]) + row_end - first_row)) + if group is not None: + calls.append((group[0], group[1], group[2], group[3], group[4])) + return calls + + +def _fp8_gemm(x: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """``x [n, K] @ weight [N, K]^T`` for a 128 x 128 block-FP8 weight: ``x``'s 1 x 128 + quantization and the blockwise FP8 GEMM, BF16 out.""" + x_fp8, x_scale = fp8_quantize_1x128(x) + return fp8_blockwise_gemm(x_fp8, weight, x_scale, scale) + + +def _linear_fp8( + x: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor, weight_bf16: torch.Tensor +) -> torch.Tensor: + """``x @ weight^T`` for a block-FP8 weight: up to ``_BF16_COPY_MAX_ROWS`` rows on + ``low_m_bf16_gemm`` over its dequantized copy ``weight_bf16``, the FP8 GEMM above.""" + if x.shape[0] <= _BF16_COPY_MAX_ROWS: + return low_m_bf16_gemm(x, weight_bf16) + return _fp8_gemm(x, weight, scale) + + +class MlaLayer: + """One sparse MLA layer's forward: the generation rows' head-parallel partial, and the + context rows' cache writes and sequence-parallel full output. + + Args: + w: the layer's weights (``load_mla``). + dims: the model's sizes and this rank's shares. + layer_idx: the layer's index (the MTP layer's is ``dims.mtp_layer``), which names its + latent and indexer pools in the KV cache manager. + aux: the aux stream a captured generation step runs its second input GEMM and its + query branch on. + scratch: the model's shared sparse attention buffers. + """ + + def __init__( + self, + w: MlaWeights, + dims: Dims, + layer_idx: int, + aux: AuxStream, + scratch: SparseAttentionScratch, + ) -> None: + if not dims.index_always_select_tail: + raise ValueError("the selection always appends the incomplete trailing pool") + if dims.kv_lora != KV_LORA_RANK: + raise ValueError( + f"the sparse attention kernel is {KV_LORA_RANK} wide, not {dims.kv_lora}" + ) + self.w = w + self.layer_idx = layer_idx + self.aux = aux + self.scratch = scratch + self.heads = dims.mla_heads + self.heads_local = dims.mla_heads_local + self.qk_nope = dims.qk_nope + self.rms_eps = dims.rms_eps + self.index_heads = dims.index_heads + self.index_head_dim = dims.index_head_dim + self.kpool = dims.index_kpool + self.select_k = dims.select_k + # select_k pools of kpool members, then the up to kpool - 1 tail positions. + width = dims.index_topk + dims.index_kpool - 1 + self.out_width = -(-width // _SELECTION_ALIGN) * _SELECTION_ALIGN + self.q_scale = dims.index_head_dim**-0.5 + self.w_scale = dims.index_heads**-0.5 + self.bmm1_scale = scratch.kv_scale_value * dims.qk_nope**-0.5 + + def generation_partial(self, x: torch.Tensor, step: StepSchedule) -> torch.Tensor: + """This rank's ``o_proj`` partial (before the TP reduction) of the step's generation rows. + + Writes the rows' latent and indexer cache rows, refreshes their pools' keys, scores and + selects pools, and attends the FP8 latent cache with this rank's heads. Graph-capturable: + shapes follow host ints only. + + Args: + x: ``[G, hidden]`` bf16, the generation rows (``G = num_generations * T``), request- + major. + step: the step's schedule. + + Returns: + ``[G, hidden]`` bf16. + """ + w, scratch = self.w, self.scratch + num_rows = x.shape[0] + assert num_rows == step.num_generation_tokens, ( + f"{num_rows} rows for {step.num_generation_tokens} generation tokens" + ) + rows = step.generation_rows() + pools = _cache_pools(step.metadata.kv_cache_manager, self.layer_idx) + tpb = pools.tokens_per_block + first = step.num_contexts + tables = step.block_tables[first : first + step.num_generations] + qkv, idx = self.aux.fork_join( + lambda: _linear_fp8(x, w.qkv_a, w.qkv_a_scale, w.qkv_a_bf16), + lambda: low_m_bf16_gemm(x, w.index_kgw), + ) + q_resid = mla_prologue( + qkv, + idx, + w.q_norm, + w.kv_norm, + w.index_k_norm, + w.index_k_bias, + self.rms_eps, + _INDEX_KEY_NORM_EPS, + rows.positions, + tables, + pools.latent, + pools.index, + tpb, + scratch.kv_scale_inv, + ) + index_weights = idx[:, 2 * self.index_head_dim :] + num_pools_max = -(-step.block_tables.shape[1] * tpb // self.kpool) + + def select() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + # The rows' pools in position order, so a pool spanning several draft rows ends with + # all of its visible members. + kpool_update( + pools.index, + tables, + rows.positions, + w.index_ape, + tpb, + head_dim=self.index_head_dim, + kpool=self.kpool, + ) + q_index = low_m_bf16_gemm(q_resid, w.index_q).view( + num_rows, self.index_heads, self.index_head_dim + ) + scores = kpool_score( + q_index, + index_weights, + pools.index, + rows.tables, + rows.visible, + tpb, + num_pools_max=num_pools_max, + head_dim=self.index_head_dim, + kpool=self.kpool, + q_scale=self.q_scale, + w_scale=self.w_scale, + ) + return kpool_select( + scores, + rows.visible, + rows.tables, + tpb, + select_k=self.select_k, + kpool=self.kpool, + base_row=pools.base_row, + rows_per_slot=pools.rows_per_slot, + out_width=self.out_width, + ) + + def query() -> torch.Tensor: + q = _linear_fp8(q_resid, w.local.q_b, w.local.q_b_scale, w.q_b_bf16) + q = q.view(num_rows, self.heads_local, self.qk_nope) + return saturate_fp8(torch.bmm(q.transpose(0, 1), w.local.w_k)) + + (packed, lengths, nonempty), query_fp8 = self.aux.fork_join(select, query) + latent = sparse_mla_decode( + query_fp8, + pools.kv_rows, + packed, + lengths, + scratch.workspace, + scratch.counter(num_rows, self.heads_local), + self.bmm1_scale, + scratch.kv_scale, + ) + # A row that selected no key (a padded request) gets zero values, so o_proj gives it zero. + values = latent_values(latent, w.local.w_v_t, nonempty) + return _linear_fp8(values, w.local.o, w.local.o_scale, w.o_bf16) + + def context_rows( + self, x: torch.Tensor, step: StepSchedule, lo: int, hi: int, out: torch.Tensor + ) -> None: + """Cache every context row, then write the full attention output of rows ``[lo, hi)``. + + Appends every context row's latent and indexer state and refreshes the keys of the pools + the step completes, as every rank does (the cache is replicated); then attends rows + ``[lo, hi)`` with all heads and writes their ``o_proj`` output, the full value rather + than a partial, into ``out[lo:hi]``. Other rows of ``out`` are not written. + + Args: + x: ``[C, hidden]`` bf16, the step's packed context rows (``C = num_ctx_tokens``). + step: the step's schedule. + lo: the first row this rank attends. + hi: one past the last; ``lo >= hi`` attends none. + out: ``[>= C, hidden]`` bf16. + """ + w, scratch = self.w, self.scratch + assert x.shape[0] == step.num_ctx_tokens, ( + f"{x.shape[0]} rows for {step.num_ctx_tokens} context tokens" + ) + sched = step.context_rows(self.kpool, x.device) + pools = _cache_pools(step.metadata.kv_cache_manager, self.layer_idx) + tpb = pools.tokens_per_block + tables = step.block_tables + qkv = _linear_fp8(x, w.qkv_a, w.qkv_a_scale, w.qkv_a_bf16) + idx = low_m_bf16_gemm(x, w.index_kgw) + q_resid = mla_context_prologue( + qkv, + idx, + w.q_norm, + w.kv_norm, + w.index_k_norm, + w.index_k_bias, + self.rms_eps, + _INDEX_KEY_NORM_EPS, + sched.positions, + sched.request_ids, + tables, + pools.latent, + pools.index, + tpb, + scratch.kv_scale_inv, + ) + del qkv + # One refresh per pool the step completes; a pool left open at the end of a chunk is + # refreshed by the write that completes it. + if sched.final_positions.shape[0]: + kpool_update( + pools.index, + tables, + sched.final_positions, + w.index_ape, + tpb, + head_dim=self.index_head_dim, + kpool=self.kpool, + request_ids=sched.final_request_ids, + ) + if lo >= hi: + return + + num_rows = hi - lo + request_ids = sched.request_ids[lo:hi] + visible = sched.positions[lo:hi] + 1 + num_cand = (visible // self.kpool).to(torch.int32) + q_resid = q_resid[lo:hi] + index_weights = idx[lo:hi, 2 * self.index_head_dim :] + q_index = low_m_bf16_gemm(q_resid, w.index_q).view( + num_rows, self.index_heads, self.index_head_dim + ) + selected = torch.empty(num_rows, self.select_k, dtype=torch.int32, device=x.device) + for start, end, first, last, max_kv_len in _score_calls( + step, lo, hi, self.kpool, self.index_head_dim + ): + part = slice(start - lo, end - lo) + ids = request_ids[part] + scores = kpool_score_fp8( + q_index[part], + index_weights[part], + pools.index, + tables[first:last], + visible[part], + tpb, + head_dim=self.index_head_dim, + kpool=self.kpool, + q_scale=self.q_scale, + w_scale=self.w_scale, + request_ids=ids - first if first else ids, + max_kv_len=max_kv_len, + ) + if scores is None: + # No complete pool is visible to these rows: they attend their tail only. + selected[part] = -1 + else: + indexer_topk(scores, num_cand[part], selected[part], select_k=self.select_k) + del scores + del q_index, index_weights, idx + topk_rows = kpool_expand( + selected, + visible, + tables, + tpb, + base_row=pools.base_row, + rows_per_slot=pools.rows_per_slot, + kpool=self.kpool, + out_width=self.out_width, + request_ids=request_ids, + ) + packed, lengths, nonempty = compact_sparse_rows(topk_rows) + del topk_rows, selected + + full = w.full + query = _fp8_gemm(q_resid, full.q_b, full.q_b_scale).view( + num_rows, self.heads, self.qk_nope + ) + query_fp8 = saturate_fp8(torch.bmm(query.transpose(0, 1), full.w_k)) + del query + latent = sparse_mla_decode( + query_fp8, + pools.kv_rows, + packed, + lengths, + scratch.workspace, + scratch.counter(num_rows, self.heads), + self.bmm1_scale, + scratch.kv_scale, + ) + del query_fp8 + values, values_scale = quantize_1x128(torch.bmm(latent.transpose(0, 1), full.w_v_t)) + del latent + projected = fp8_blockwise_gemm(values, full.o, values_scale, full.o_scale) + # A row that selected no key has an unspecified attention output; it projects to zero. + out[lo:hi] = torch.where(nonempty[:, None], projected, 0) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/layers/streams.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/layers/streams.py new file mode 100644 index 000000000000..2dbe85d1e8d1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/layers/streams.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The aux stream two independent pieces of a captured generation step run on side by side.""" + +from __future__ import annotations + +from typing import Callable, TypeVar + +import torch + +_Main = TypeVar("_Main") +_Aux = TypeVar("_Aux") + + +class AuxStream: + """One aux stream and the fork and join events that order it against the current stream. + + Made once per model after the weights load: the engine constructs the model on the meta + device, and a stream made inside a CUDA graph capture is not capturable. + """ + + def __init__(self) -> None: + self.stream = torch.cuda.Stream() + self.fork = torch.cuda.Event() + self.join = torch.cuda.Event() + + def fork_join(self, main: Callable[[], _Main], aux: Callable[[], _Aux]) -> tuple[_Main, _Aux]: + """``(main(), aux())``: inside a CUDA graph capture ``aux`` runs on the aux stream beside + ``main``; otherwise both run in order on the current stream.""" + if not torch.cuda.is_current_stream_capturing(): + return main(), aux() + # In a capture the aux stream's outputs need no record_stream: every use of the aux + # stream starts by waiting for the current stream (the fork). + self.fork.record() + main_out = main() + with torch.cuda.stream(self.stream): + self.fork.wait() + aux_out = aux() + self.join.record() + self.join.wait() + return main_out, aux_out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/modeling.py new file mode 100644 index 000000000000..1ad2fbeb98e4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/modeling.py @@ -0,0 +1,839 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GLM-5.3-Flash block-FP8 on sm_100, tp4: model construction and forward. + +The checkpoint is the composite ``Glm5NextForConditionalGeneration`` release, served text +only (``disable_mm_encoder``). The engine builds ``ModelingV2Glm53FlashFp8Sm10xTp4``, a +``DecoderModelForCausalLM`` shell around ``Glm53FlashFp8Sm10xTp4Core``; one-model MTP-3 +drafts with the checkpoint's MTP layer through the engine's speculative worker. + +Each step runs one of two targets: + +- ``PrefillTarget`` for steps with context rows, sequence parallel: rank ``r`` owns rows + ``[r * n, (r + 1) * n)`` of the packed step (``n = ceil(T / 4)``) and keeps the mHC + streams of those rows only, so every boundary (post-mapping, pre-mapping, RMSNorm) runs + over a quarter of the rows. Each sublayer is all-gather, the sublayer over all rows, + reduce-scatter: the same bytes as the all-reduce it replaces. KDA, the MoE and the dense + MLPs produce head- or intermediate-split partials as under TP. The MLA layers attend + their context rows sequence parallel (all 64 heads, the owned rows only, zero for the + others, so the reduce-scatter leaves each owned row's full value); generation rows of a + mixed step stay head-parallel. +- ``DecodeTarget`` for generation-only steps (the steps CUDA graphs capture): the decoder + over every row under TP, with the TP all-reduce of a step of at most four rows fused into + the mHC boundary after each sublayer. + +Every layer's computation is this target's: ``layers/`` holds the KDA, sparse MLA, FFN and +mHC layers on catalog and target operations, ``weights.py`` reads the checkpoint, and +``schedule.py`` derives the step's row schedules from the engine's prepared metadata. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Mapping, Optional, Union + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tensorrt_llm._ipc_utils import can_access_peer +from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + disable_v2_by_default, + step_contract_enabled, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.allgather import allgather +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.allreduce import ( + allreduce, + allreduce_group_args, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.reducescatter import reducescatter +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata +from tensorrt_llm._torch.attention.backends.sparse.glm_kpool.metadata import Glm5NextMamba2Metadata +from tensorrt_llm._torch.distributed.ops import get_allreduce_workspace +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_utils import DecoderModel, DecoderModelForCausalLM +from tensorrt_llm._torch.speculative import get_spec_worker, should_use_separate_draft_kv_cache +from tensorrt_llm._utils import mpi_disabled +from tensorrt_llm.functional import AllReduceFusionOp, AllReduceStrategy +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping as ParallelMapping + +from ..routing import fp8_model_refused +from . import weights as _weights +from .dims import Dims +from .layers.ffn import DenseMlp, MoeLayer +from .layers.kda import KdaLayer +from .layers.mhc import Mhc +from .layers.mla import MlaLayer, SparseAttentionScratch +from .layers.streams import AuxStream +from .ops.comm.mhc_ar_boundary import WORLD as _FUSED_BOUNDARY_WORLD +from .ops.comm.mhc_ar_boundary import mhc_ar_boundary +from .ops.gemm.low_m_bf16_gemm import low_m_bf16_gemm +from .ops.moe.decode_moe import MAX_ROWS as _DECODE_MOE_ROWS +from .ops.moe.prefill_moe import PrefillMoeWorkspace +from .ops.quantization.fp8_quantize_1x128 import fp8_quantize_1x128 +from .schedule import DENSE_MLP, LINEAR_ATTENTION, StepSchedule + +if TYPE_CHECKING: + from tensorrt_llm._torch.pyexecutor.guided_decoder import CapturableGuidedDecoder + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest + from tensorrt_llm._torch.speculative.interface import SpecMetadata + +#: The TP all-reduce of more rows than this runs on NCCL; fewer on the one-shot kernel, +#: whose workspace the fused boundary shares. +_ONESHOT_MAX_ROWS = 512 +_AR_NONE = int(AllReduceFusionOp.NONE) + +#: The prepared hybrid-cache metadata fields the KDA layers read (``StepSchedule.build`` reads +#: the GLM ones). +_HYBRID_FIELDS = ( + "state_indices", + "generation_state_indices", + "has_initial_states", + "query_start_loc", + "query_start_loc_long", + "kda_max_seqlen", + "kda_seq_order", + "kda_cu_chunks", +) +#: The KV cache manager members the layers use: the MLA layers' latent and indexer pools and +#: the KDA layers' state. +_CACHE_MANAGER_MEMBERS = ( + "tokens_per_block", + "get_latent_state_buffer", + "get_index_state_buffer", + "mamba_layer_cache", +) +#: A KDA layer cache's fused-verify replay members, read with MTP. +_KDA_REPLAY_FIELDS = ( + "kda_conv_q", + "kda_conv_k", + "kda_conv_v", + "kda_k_cache", + "kda_g_cache", + "kda_v_cache", + "kda_beta_cache", + "prev_num_accepted_tokens", + "commit_conv_window", +) + + +class _Comm: + """This rank's TP collectives: the all-reduce (one-shot up to ``_ONESHOT_MAX_ROWS`` rows, + NCCL above), and the all-gather / reduce-scatter of sequence-parallel prefill.""" + + def __init__(self, mapping: ParallelMapping) -> None: + self.tp_size = mapping.tp_size + self.tp_rank = mapping.tp_rank + self.group = list(mapping.tp_group) + self.process_group = mapping.tp_group_pg.boxed() if mpi_disabled() else None + self.small = allreduce_group_args(mapping, int(AllReduceStrategy.ONESHOT)) + self.large = allreduce_group_args(mapping, int(AllReduceStrategy.NCCL)) + # The one-shot all-reduce's Lamport workspace, which the fused boundary also rotates. + self.workspace = get_allreduce_workspace(mapping) + # The fused boundary stores through the workspace's peer pointers, which are null + # without peer access (the library all-reduce then falls back to NCCL). + self.peer_access = can_access_peer(mapping) + + def allreduce(self, x: torch.Tensor) -> torch.Tensor: + args = self.small if x.shape[0] <= _ONESHOT_MAX_ROWS else self.large + return allreduce( + x.contiguous(), + None, + None, + None, + None, + op=_AR_NONE, + eps=0.0, + trigger_completion_at_end=False, + **args, + ) + + def gather(self, local: torch.Tensor, num_rows: int) -> torch.Tensor: + """Every rank's rows ``local`` stacked in rank order, cut to ``num_rows``.""" + return allgather(local, None, self.group, self.process_group)[:num_rows] + + def scatter(self, partial: torch.Tensor, rows: int) -> torch.Tensor: + """This rank's ``rows`` rows of the sum of every rank's ``partial``.""" + pad = rows * self.tp_size - partial.shape[0] + if pad: + partial = torch.nn.functional.pad(partial, (0, 0, 0, pad)) + return reducescatter(partial, None, self.group, self.process_group) + + def gather_vocab(self, local: torch.Tensor) -> torch.Tensor: + """``[rows, vocab / tp]`` shards -> ``[rows, vocab]``, rank-major along the vocabulary.""" + rows = local.shape[0] + stacked = allgather(local.contiguous(), None, self.group, self.process_group) + return stacked.view(self.tp_size, rows, -1).transpose(0, 1).reshape(rows, -1) + + +def _expand_streams(x: torch.Tensor, mult: int) -> torch.Tensor: + """``[rows, hidden]`` -> ``[rows, mult, hidden]`` exact copies: the streams diverge at the + first mixer. Contiguous, because post-mapping writes each stream separately.""" + return x.unsqueeze(-2).expand(-1, mult, -1).contiguous() + + +def _collapse_streams(streams: torch.Tensor, norm_weight: torch.Tensor, eps: float) -> torch.Tensor: + """The decoder's output: the streams' unweighted mean (GLM-5.3-Flash's stack ends in a plain + mean, not a learned head), then the final RMSNorm.""" + return flashinfer_rmsnorm(streams.mean(dim=-2), norm_weight, eps) + + +class _Layer: + """One decoder layer's parts: the attention (KDA or MLA), the FFN, the two mHC sites and + the two RMSNorm weights.""" + + def __init__( + self, + index: int, + kind: str, + attention: Union[KdaLayer, MlaLayer], + mlp: Union[DenseMlp, MoeLayer], + hc_attn: Mhc, + hc_ffn: Mhc, + input_norm: torch.Tensor, + post_norm: torch.Tensor, + ) -> None: + self.index = index + self.kind = kind + self.attention = attention + self.mlp = mlp + self.hc_attn = hc_attn + self.hc_ffn = hc_ffn + self.input_norm = input_norm + self.post_norm = post_norm + + @property + def is_kda(self) -> bool: + return self.kind == LINEAR_ATTENTION + + +class _Mtp: + """The MTP layer: front end, sparse MLA, MoE, head norm; plain residuals, no mHC.""" + + def __init__(self, w: _weights.MtpLayerWeights, attention: MlaLayer, mlp: MoeLayer) -> None: + self.w = w + self.attention = attention + self.mlp = mlp + + +class PrefillTarget: + """Steps with context rows, sequence parallel over the step's rows (see the module + docstring).""" + + def __init__(self, core: "Glm53FlashFp8Sm10xTp4Core") -> None: + self.core = core + self.comm = core.comm + self.moe_workspace = PrefillMoeWorkspace() + + def forward( + self, step: StepSchedule, input_ids: torch.Tensor, attn_metadata: AttentionMetadata + ) -> torch.Tensor: + core = self.core + comm = self.comm + dims = core.dims + eps = dims.rms_eps + embeds = torch.nn.functional.embedding(input_ids, core.w.embed) + num_tokens = embeds.shape[0] + rows = -(-num_tokens // comm.tp_size) + first_row = comm.tp_rank * rows + own = embeds[first_row : first_row + rows] + if own.shape[0] < rows: + own = torch.nn.functional.pad(own, (0, 0, 0, rows - own.shape[0])) + + layers = core.layers + residual = _expand_streams(own, dims.hc_mult) + post, comb, x = layers[0].hc_attn.pre_mapping(residual) + x = flashinfer_rmsnorm(x, layers[0].input_norm, eps) + for i, layer in enumerate(layers): + full = comm.gather(x, num_tokens) + if layer.is_kda: + partial = layer.attention.partial(full, step, attn_metadata) + else: + partial = self._mla_partial(layer.attention, full, step, rows) + del full + residual, post, comb, x = layer.hc_ffn.fused_hc( + comm.scatter(partial, rows), residual, post, comb, layer.post_norm, eps + ) + del partial + mlp_out = comm.scatter(self._ffn_partial(layer.mlp, x, num_tokens), rows) + if i + 1 < len(layers): + nxt = layers[i + 1] + residual, post, comb, x = nxt.hc_attn.fused_hc( + mlp_out, residual, post, comb, nxt.input_norm, eps + ) + else: + residual = layer.hc_ffn.post_mapping(mlp_out, residual, post, comb) + del mlp_out + return comm.gather(_collapse_streams(residual, core.w.final_norm, eps), num_tokens) + + def mtp_forward( + self, step: StepSchedule, input_ids: torch.Tensor, hidden_states: torch.Tensor + ) -> torch.Tensor: + """The MTP layer over a step with context rows, sequence parallel like the decoder: + its row-parallel eh_proj reduce-scatters to the owned rows, the attention and the FFN + are all-gather, sublayer, reduce-scatter, and the head-normalized output is gathered + back to every row.""" + core = self.core + comm = self.comm + mtp = core.mtp + eps = core.dims.rms_eps + num_tokens = hidden_states.shape[0] + rows = -(-num_tokens // comm.tp_size) + own_cols = _mtp_own_columns(core, input_ids, hidden_states) + residual = comm.scatter(torch.nn.functional.linear(own_cols, mtp.w.eh_proj), rows) + full = comm.gather(flashinfer_rmsnorm(residual, mtp.w.input_norm, eps), num_tokens) + partial = self._mla_partial(mtp.attention, full, step, rows) + del full + residual = residual + comm.scatter(partial, rows) + x = flashinfer_rmsnorm(residual, mtp.w.post_norm, eps) + residual = residual + comm.scatter(self._ffn_partial(mtp.mlp, x, num_tokens), rows) + return comm.gather(flashinfer_rmsnorm(residual, mtp.w.head_norm, eps), num_tokens) + + def _ffn_partial( + self, mlp: Union[DenseMlp, MoeLayer], x: torch.Tensor, num_tokens: int + ) -> torch.Tensor: + """The FFN's intermediate-split output for every row, from this rank's rows ``x``. + + Every FFN GEMM input is the 1 x 128 block-FP8 quantization of the row, a per-row + function: each rank quantizes its own rows (and runs the router on them), and the + gather moves FP8 rows, their scales and the router logits instead of BF16 rows every + rank would quantize again. + """ + comm = self.comm + x_fp8, x_scale = fp8_quantize_1x128(x) + x_fp8 = comm.gather(x_fp8.view(torch.uint8), num_tokens).view(torch.float8_e4m3fn) + # Gathered row-major, then returned to the K-block-major layout the GEMMs read. + x_scale = comm.gather(x_scale.t().contiguous(), num_tokens).t().contiguous() + if isinstance(mlp, MoeLayer): + logits = comm.gather(mlp.router_logits(x), num_tokens) + return mlp.prefill_partial(x_fp8, x_scale, logits, self.moe_workspace) + return mlp.partial_from_fp8(x_fp8, x_scale) + + def _mla_partial( + self, attention: MlaLayer, x: torch.Tensor, step: StepSchedule, rows: int + ) -> torch.Tensor: + """This rank's contribution to the MLA output of every row: the full value of its own + context rows, zero for the other context rows, and the head-shard partial of the + generation rows.""" + out = x.new_zeros(x.shape) + num_ctx = step.num_ctx_tokens + if step.num_contexts: + lo = min(self.comm.tp_rank * rows, num_ctx) + hi = min(lo + rows, num_ctx) + attention.context_rows(x[:num_ctx], step, lo, hi, out) + if step.num_generations: + out[num_ctx:] = attention.generation_partial(x[num_ctx:], step) + return out + + +class DecodeTarget: + """Generation-only steps under TP: every row is a generation (or verify) row. A step of at + most ``_DECODE_MOE_ROWS`` rows ends each sublayer unreduced and fuses the TP all-reduce + into the mHC boundary after it.""" + + def __init__(self, core: "Glm53FlashFp8Sm10xTp4Core") -> None: + self.core = core + self.comm = core.comm + + def forward( + self, step: StepSchedule, input_ids: torch.Tensor, attn_metadata: AttentionMetadata + ) -> torch.Tensor: + core = self.core + comm = self.comm + eps = core.dims.rms_eps + layers = core.layers + x = torch.nn.functional.embedding(input_ids, core.w.embed) + fused = ( + x.shape[0] <= _DECODE_MOE_ROWS + and comm.tp_size == _FUSED_BOUNDARY_WORLD + and comm.peer_access + ) + residual = _expand_streams(x, core.dims.hc_mult) + post, comb, x = layers[0].hc_attn.pre_mapping(residual) + x = flashinfer_rmsnorm(x, layers[0].input_norm, eps) + for i, layer in enumerate(layers): + if layer.is_kda: + partial = layer.attention.partial(x, step, attn_metadata) + else: + partial = layer.attention.generation_partial(x, step) + if fused: + residual, post, comb, x = self._fused_boundary( + partial, residual, post, comb, layer.hc_ffn, layer.post_norm + ) + else: + residual, post, comb, x = layer.hc_ffn.fused_hc( + comm.allreduce(partial), residual, post, comb, layer.post_norm, eps + ) + del partial + last = i + 1 == len(layers) + if fused and not last and isinstance(layer.mlp, MoeLayer): + nxt = layers[i + 1] + residual, post, comb, x = self._fused_boundary( + layer.mlp.partial(x), residual, post, comb, nxt.hc_attn, nxt.input_norm + ) + continue + mlp_out = comm.allreduce(layer.mlp.partial(x)) + if not last: + nxt = layers[i + 1] + residual, post, comb, x = nxt.hc_attn.fused_hc( + mlp_out, residual, post, comb, nxt.input_norm, eps + ) + else: + residual = layer.hc_ffn.post_mapping(mlp_out, residual, post, comb) + del mlp_out + return _collapse_streams(residual, core.w.final_norm, eps) + + def mtp_forward( + self, step: StepSchedule, input_ids: torch.Tensor, hidden_states: torch.Tensor + ) -> torch.Tensor: + """The MTP layer over generation rows under TP: eh_proj's row-parallel partial, then + attention and FFN with plain residuals, each reduced once.""" + core = self.core + comm = self.comm + mtp = core.mtp + eps = core.dims.rms_eps + own_cols = _mtp_own_columns(core, input_ids, hidden_states) + residual = comm.allreduce(_linear(own_cols, mtp.w.eh_proj)) + x = flashinfer_rmsnorm(residual, mtp.w.input_norm, eps) + residual = residual + comm.allreduce(mtp.attention.generation_partial(x, step)) + x = flashinfer_rmsnorm(residual, mtp.w.post_norm, eps) + residual = residual + comm.allreduce(mtp.mlp.partial(x)) + return flashinfer_rmsnorm(residual, mtp.w.head_norm, eps) + + def _fused_boundary( + self, + partial: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + hc: Mhc, + norm_weight: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """What ``hc.fused_hc(allreduce(partial), residual, post, comb, norm_weight, eps)`` + returns, with the TP all-reduce of ``partial`` fused in.""" + rows, mult = residual.shape[0], hc.mult + residual_out, post_out, comb_out, layer_input = mhc_ar_boundary( + partial, + self.comm.workspace, + self.comm.tp_rank, + residual.reshape(rows, mult, -1), + post.reshape(rows, mult), + comb.reshape(rows, mult, mult), + hc.fn, + hc.scale, + hc.base, + norm_weight, + *hc.scalars(self.core.dims.rms_eps), + ) + return residual_out.view_as(residual), post_out.view(rows, mult), comb_out, layer_input + + +def _linear(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """``x @ weight.T`` for a BF16 weight: ``low_m_bf16_gemm``, which takes its low-M kernels + for the shapes and row counts it covers and cuBLAS otherwise.""" + rows = x.reshape(-1, x.shape[-1]).contiguous() + return low_m_bf16_gemm(rows, weight).view(*x.shape[:-1], -1) + + +def _mtp_own_columns( + core: "Glm53FlashFp8Sm10xTp4Core", input_ids: torch.Tensor, hidden_states: torch.Tensor +) -> torch.Tensor: + """This rank's input columns of the MTP layer's row-parallel eh_proj: + ``cat(enorm(embed(input_ids)), hnorm(hidden_states))[:, rank's quarter]``.""" + mtp = core.mtp + eps = core.dims.rms_eps + e = flashinfer_rmsnorm(torch.nn.functional.embedding(input_ids, core.w.embed), mtp.w.enorm, eps) + h = flashinfer_rmsnorm(hidden_states, mtp.w.hnorm, eps) + joined = torch.cat([e, h], dim=-1) + tp, rank = core.comm.tp_size, core.comm.tp_rank + return torch.chunk(joined, tp, dim=-1)[rank].contiguous() + + +class Glm53FlashFp8Sm10xTp4Core(DecoderModel): + """The text decoder: this rank's weights, the layers built on them, and the per-step choice + between the two targets.""" + + def __init__(self, model_config: ModelConfig, text_config: PretrainedConfig) -> None: + super().__init__(model_config) + spec = model_config.spec_config + mtp = spec is not None and spec.spec_dec_mode.is_mtp_one_model() + self.text_config = text_config + self.dims = Dims.from_config( + text_config, + model_config.mapping, + mtp, + model_config.moe_max_num_tokens or model_config.max_num_tokens, + ) + self.w: Optional[_weights.ModelWeights] = None + self.layers: list[_Layer] = [] + self.mtp: Optional[_Mtp] = None + self.comm: Optional[_Comm] = None + self._prefill: Optional[PrefillTarget] = None + self._decode: Optional[DecodeTarget] = None + # Read once: TRTLLM_MODELING_V2_VALIDATE asks the first forward to check the metadata + # surface this target reads (``_check_step_contract``). + self._contract_pending = step_contract_enabled() + + def load(self, weights: Mapping[str, Any], lm_head_weight: torch.Tensor) -> None: + device = torch.device("cuda", torch.cuda.current_device()) + self.w = _weights.load_model_weights(weights, self.dims, lm_head_weight, device) + + @torch.no_grad() + def post_load_weights(self) -> None: + """Build the layers and the two targets once the weights are real. Reached more than + once (the shell and the loader's module walk), so it builds once. Streams and the + collectives' workspaces are made here, not at construction: the engine constructs the + model on the meta device, and a stream made inside a graph capture is not capturable.""" + if self._decode is not None: + return + assert self.w is not None, "post_load_weights before load_weights" + dims = self.dims + mapping = self.model_config.mapping + self.comm = _Comm(mapping) + aux = AuxStream() + scratch = SparseAttentionScratch(dims) + self.layers = [] + for i, w in enumerate(self.w.layers): + if w.attention_kind == LINEAR_ATTENTION: + attention = KdaLayer(w.attention, dims, i, aux) + else: + attention = MlaLayer(w.attention, dims, i, aux, scratch) + mlp = DenseMlp(w.mlp, dims, i) if w.mlp_kind == DENSE_MLP else MoeLayer(w.mlp, dims, i) + self.layers.append( + _Layer( + i, + w.attention_kind, + attention, + mlp, + Mhc(w.hc_attn, dims), + Mhc(w.hc_ffn, dims), + w.input_norm, + w.post_norm, + ) + ) + if self.w.mtp is not None: + m = self.w.mtp + layer = dims.mtp_layer + self.mtp = _Mtp( + m, + MlaLayer(m.attention, dims, layer, aux, scratch), + MoeLayer(m.mlp, dims, layer), + ) + self._prefill = PrefillTarget(self) + self._decode = DecodeTarget(self) + + def select_target(self, attn_metadata: AttentionMetadata) -> "PrefillTarget | DecodeTarget": + """A step with any context row is the prefill target's (mixed steps included); a + generation-only step is the decode target's. One host int read, so a captured decode + graph sees a constant.""" + if self._decode is None: + self.post_load_weights() + return self._decode if attn_metadata.num_contexts == 0 else self._prefill + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + **kwargs: Any, + ) -> torch.Tensor: + assert inputs_embeds is None, "this text-only target embeds the input ids itself" + target = self.select_target(attn_metadata) + if self._contract_pending: + self._check_step_contract(attn_metadata) + return target.forward(StepSchedule.build(attn_metadata), input_ids, attn_metadata) + + def _check_step_contract(self, attn_metadata: AttentionMetadata) -> None: + """Opt-in first-forward fail-fast (``TRTLLM_MODELING_V2_VALIDATE``): every field of the + engine-prepared metadata and every KV cache manager member the layers read must exist, + so an upstream rename fails here rather than mid-forward. They are fixed at engine + construction, so one check per instance is sound.""" + try: + StepSchedule.build(attn_metadata) + for name in ("seq_lens_cuda", "padded_num_tokens"): + getattr(attn_metadata, name) + for name in _HYBRID_FIELDS: + getattr(attn_metadata.mamba_metadata, name) + manager = attn_metadata.kv_cache_manager + for name in _CACHE_MANAGER_MEMBERS: + getattr(manager, name) + cache = manager.mamba_layer_cache(next(x.index for x in self.layers if x.is_kda)) + # A disaggregated context server's manager allocates no verify replay caches. + verifies = self.dims.mtp and not getattr(manager, "kda_context_only", False) + for name in ("conv", "temporal") + (_KDA_REPLAY_FIELDS if verifies else ()): + getattr(cache, name) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + def mtp_forward( + self, + input_ids: torch.Tensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + """The MTP layer on the metadata of this draft call: the worker rewrites the metadata + in place between draft steps, so the target is chosen per call.""" + step = StepSchedule.build(attn_metadata) + return self.select_target(attn_metadata).mtp_forward(step, input_ids, hidden_states) + + +class _VocabParallelHead: + """The shell's vocabulary-parallel LM head on this target's GEMMs: decode-sized products on + ``low_m_bf16_gemm``, the rest on cuBLAS, this rank's vocabulary shard or, with + ``gather_output``, every rank's shards rank-major. The speculative worker and the logits + processor call it as they call the shell's head, whose TP attributes it carries.""" + + def __init__(self, lm_head: nn.Module, comm: _Comm) -> None: + self.lm_head = lm_head + self.comm = comm + self.tp_mode = lm_head.tp_mode + self.gather_output = lm_head.gather_output + self.padding_size = getattr(lm_head, "padding_size", 0) + + @property + def weight(self) -> torch.Tensor: + return self.lm_head.weight + + def __call__( + self, + hidden_states: torch.Tensor, + gather_output: Optional[bool] = None, + is_spec_decoding_head: bool = False, + **kwargs: Any, + ) -> torch.Tensor: + local = _linear(hidden_states, self.lm_head.weight) + gather = self.gather_output if gather_output is None else gather_output + if not gather: + return local + logits = self.comm.gather_vocab(local) + # The last rank's shard is padded to the shard size; the padding is no vocabulary entry. + return logits[..., : -self.padding_size] if self.padding_size else logits + + +class MTPLayer: + """The checkpoint's MTP layer as the speculative worker drives it: ``__call__`` runs the + layer (whose output is already ``shared_head.norm``-ed, unlike DeepSeek's) and + ``shared_head`` projects it, without another norm, to this rank's vocabulary shard of + the draft logits; the worker's greedy draft sampler gathers the argmax across ranks.""" + + def __init__(self, core: Glm53FlashFp8Sm10xTp4Core) -> None: + self.core = core + + def __call__( + self, + embed_tokens: Optional[torch.Tensor] = None, + all_rank_num_tokens: Optional[list[int]] = None, + input_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + hidden_states: Optional[torch.Tensor] = None, + attn_metadata: Optional[AttentionMetadata] = None, + spec_metadata: Optional[SpecMetadata] = None, + **kwargs: Any, + ) -> torch.Tensor: + del embed_tokens, all_rank_num_tokens, position_ids, kwargs # NoPE; own embedding + hidden_states = self.core.mtp_forward(input_ids, hidden_states, attn_metadata) + if spec_metadata is not None: + spec_metadata.maybe_capture_hidden_states(0, hidden_states, None) + return hidden_states + + def shared_head( + self, + hidden_states: torch.Tensor, + lm_head: nn.Module, + attn_metadata: AttentionMetadata, + return_context_logits: bool = False, + ) -> torch.Tensor: + if not return_context_logits: + last = torch.cumsum(attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1 + hidden_states = hidden_states[last] + return _linear(hidden_states, lm_head.weight) + + +class DraftModel: + """The container the speculative worker reaches the drafter through. ``embed_tokens`` + resolves through the core on every read: the container is built before the weights are.""" + + def __init__(self, core: Glm53FlashFp8Sm10xTp4Core, lm_head: nn.Module) -> None: + self.core = core + self.mtp_layers = [MTPLayer(core)] + self.lm_head = lm_head + + @property + def embed_tokens(self) -> torch.Tensor: + return self.core.w.embed + + +def _text_config(pretrained_config: PretrainedConfig) -> PretrainedConfig: + return getattr(pretrained_config, "text_config", None) or pretrained_config + + +def _normalize_top_config(config: PretrainedConfig, text: PretrainedConfig) -> None: + """Expose on the composite config the text fields the engine reads off it (capacity + planning, the speculative metadata and the LM head), keeping values already there.""" + for name in ("num_hidden_layers", "hidden_size", "vocab_size"): + if getattr(config, name, None) is None: + setattr(config, name, int(getattr(text, name))) + if getattr(config, "torch_dtype", None) is None: + config.torch_dtype = getattr(text, "torch_dtype", None) or torch.bfloat16 + if getattr(config, "num_nextn_predict_layers", None) is None: + config.num_nextn_predict_layers = int(getattr(text, "num_nextn_predict_layers", 0) or 0) + + +@disable_v2_by_default +class ModelingV2Glm53FlashFp8Sm10xTp4( + DecoderModelForCausalLM[Glm53FlashFp8Sm10xTp4Core, PretrainedConfig] +): + """GLM-5.3-Flash block-FP8, sm_100, tp4, text only.""" + + #: The KDA layers reset a fresh request's recurrent state and conv window in their prefill + #: only, so the engine must not run a one-token prompt as a decode row. + resets_recurrent_state_in_context = True + + def __init__(self, model_config: ModelConfig, *args: Any, **kwargs: Any) -> None: + # ``route`` declines the same facts; this covers a direct construction. + refused = fp8_model_refused(model_config, torch.cuda.get_device_capability()) + if not model_config.disable_mm_encoder: + refused.append("the vision encoder (only disable_mm_encoder)") + if refused: + raise ValueError(f"the block-FP8 glm5_next target does not run {'; '.join(refused)}") + mapping = model_config.mapping + config = model_config.pretrained_config + text = _text_config(config) + _normalize_top_config(config, text) + super().__init__( + Glm53FlashFp8Sm10xTp4Core(model_config, text), + config=model_config, + hidden_size=int(text.hidden_size), + vocab_size=int(text.vocab_size), + ) + self.text_config = text + # device='cpu' keeps the constant real when the engine constructs on the meta device. + self._mm_token_ids = torch.tensor( + [int(config.image_token_id), int(config.video_token_id)], + dtype=torch.int32, + device="cpu", + ) + self.spec_config = model_config.spec_config + self.draft_model = None + self.spec_worker = None + self._head: Optional[_VocabParallelHead] = None + if self.model.dims.mtp: + self.draft_model = DraftModel(self.model, self.lm_head) + self.spec_worker = get_spec_worker( + self.spec_config, + model_config, + mapping, + use_separate_draft_kv_cache=should_use_separate_draft_kv_cache(self.spec_config), + ) + self.spec_worker.set_draft_model(self.draft_model) + self.spec_worker.set_target_model(self) + logger.info(f"modeling_v2: {type(self).__name__} built") + + # -- engine contracts ---------------------------------------------------------------- + + @property + def mamba_metadata_cls(self) -> type[Glm5NextMamba2Metadata]: + """The hybrid-cache metadata the engine prepares for this model: GLM block tables, KDA + state indices and the context schedules ``schedule.StepSchedule`` reads.""" + return Glm5NextMamba2Metadata + + @property + def multimodal_token_ids(self) -> torch.Tensor: + """The image and video placeholder ids, which this text-only target never embeds.""" + return self._mm_token_ids + + def infer_max_seq_len(self) -> int: + """The text config's context limit; the model has no RoPE (NoPE MLA, KDA).""" + return int(self.text_config.max_position_embeddings) + + def validate_request(self, request: LlmRequest) -> None: + """Reject media on admission, as that request's own error: the vision tower is not built, + and the forward would assert on the whole step.""" + data = request.py_multimodal_data or {} + media = [kind for kind in ("image", "video") if data.get(kind)] + if media: + raise ValueError( + f"this GLM-5.3-Flash deployment serves text only (disable_mm_encoder); the " + f"request carries {' and '.join(media)} input" + ) + + def set_guided_decoder(self, guided_decoder: CapturableGuidedDecoder) -> bool: + if self.spec_worker is not None: + return self.spec_worker.set_guided_decoder(guided_decoder) + return False + + # -- weights ------------------------------------------------------------------------- + + def load_weights(self, weights: Mapping[str, Any], *args: Any, **kwargs: Any) -> None: + assert self.lm_head.weight.dtype == torch.bfloat16, ( + f"the checkpoint's lm_head is BF16, the shell built {self.lm_head.weight.dtype}" + ) + self.model.load(weights, self.lm_head.weight) + + def post_load_weights(self) -> None: + super().post_load_weights() + self.model.post_load_weights() + if self._head is None: + self._head = _VocabParallelHead(self.lm_head, self.model.comm) + + # -- forward ------------------------------------------------------------------------- + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[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, + lora_params: Optional[dict] = None, + **kwargs: Any, + ) -> Union[torch.Tensor, dict[str, torch.Tensor]]: + """One engine step. ``resource_manager`` stays inside ``**kwargs``, and ``position_ids`` + reaches the worker in the engine's ``[1, T]`` shape (the worker squeezes it itself). + + A step padded for a prefill CUDA graph or attention-DP peers (``padded_num_tokens``, + neither of which this deployment runs) is cut back to its ``num_tokens`` rows before + the decoder, as the built-in model cuts its output. Context logits go through the + shell's ``LMHead``, whose bounded context-logits plan the logits processor applies.""" + multimodal = kwargs.get("multimodal_params") + assert not multimodal, "this target serves text only" + if attn_metadata.padded_num_tokens is not None: + num_tokens = attn_metadata.num_tokens + input_ids = input_ids[:num_tokens] + if position_ids is not None: + position_ids = position_ids[..., :num_tokens] + hidden = self.model( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + ) + head = self._head + if self.spec_worker is None: + if return_context_logits: + head = self.lm_head + return self.logits_processor.forward(hidden, head, attn_metadata, return_context_logits) + logits = self.spec_worker.target_logits( + hidden[spec_metadata.gather_ids], + head, + self.logits_processor, + attn_metadata, + spec_metadata, + self.draft_model, + ) + return self.spec_worker( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden, + logits=logits, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=self.draft_model, + resource_manager=kwargs.get("resource_manager"), + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/__init__.py new file mode 100644 index 000000000000..f8bbe8b7c789 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/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/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/activation/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/activation/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/activation/__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/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/activation/silu_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/activation/silu_and_mul.py new file mode 100644 index 000000000000..fd1e56767443 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/activation/silu_and_mul.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gated SiLU over a fused `[gate | up]` row, with the optional SwiGLU clamp, via the trtllm +Triton silu_and_mul kernel.""" + +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 + +#: GLM-5.3-Flash's `swiglu_limit`. +GLM_SWIGLU_LIMIT = 10.0 + + +class SiluAndMul(OpWrapper): + """`silu(gate) * up` over `x = [gate | up]` (`d = x.shape[1] // 2`), new `[M, d]` tensor out. + + The gate is the FIRST half and the up projection the second -- the order the fused + `gate_up` weight of a dense MLP or shared expert stores (`[gate; up]` rows). The routed + experts of `fp8_block_scale_moe_runner` store the opposite order (`[up; gate]`), so the two + halves are not interchangeable between those weights. + + With `swiglu_limit = L > 0`, before the product: `gate = min(gate, L)` (upper side only) + and `up = clamp(up, -L, L)`. Then `out = gate * sigmoid(alpha * gate) * (up + beta)` in + fp32, with `alpha = 1`, `beta = 0` unless given, rounded once to `dtype` (default + `x.dtype`). With `scale` (and an fp8 `dtype`) the result is divided by `scale[0]` and + saturated to the dtype's finite range before the cast. + + Fusion boundary: the clamp, SiLU, product and optional output quantization; no GEMM. + """ + + ARCHS = frozenset({Arch.SM_100}) + + # Only GLM-5.3-Flash calls this entry: its dense MLP and shared expert, with the clamp on. + # Its own cells are the kernel's tiling edges at those widths; the target's row counts are + # in the test. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one row of the shared expert at TP4: d = 512 is narrower than the kernel's " + "1024-element block, so every program is masked", + spec=dict(rows=1, cols=1024, swiglu_limit=GLM_SWIGLU_LIMIT), + ), + Cell( + why="one row of the dense MLP at TP4: d = 3072 spans three 1024-element blocks", + spec=dict(rows=1, cols=6144, swiglu_limit=GLM_SWIGLU_LIMIT), + ), + ) + + note = """ + Source: a Triton kernel (`silu_and_mul_kernel` in `_torch/modules/swiglu.py`, launched by + the `trtllm::silu_and_mul` custom op in `_torch/custom_ops/torch_custom_ops.py`), one + program per row and 1024-element column block; JIT-compiled per distinct `swiglu_limit`, + `swiglu_alpha` and `swiglu_beta` (all three are compile-time constants). + + The clamp is compiled in only for `swiglu_limit > 0`: `None`, `0.0` and negative limits + all mean no clamp. That differs from `fp8_block_scale_moe_runner`, whose clamp is applied + whenever a limit is given. + + Preconditions the op enforces loudly: `x` 2-D with an even width. Each row is read from + `x.stride(0)`, so a row-sliced view of a wider buffer is fine; the columns are read as + dense (see `is_valid`). The output is a new contiguous tensor. + + Certified: bf16 in and out, `swiglu_limit` 10.0 (GLM-5.3-Flash), `alpha`/`beta` unset, + no `scale`, at GLM's per-rank widths (2 x 512 and 2 x 3072) from 1 to 13,312 rows. + + Not certified: an fp8 output (`scale` + `dtype`), `swiglu_alpha` / `swiglu_beta`, other + input dtypes, CUDA-graph capture, sm_103. + """ + + def __call__( + self, + x: torch.Tensor, + scale: Optional[torch.Tensor] = None, + dtype: Optional[torch.dtype] = None, + swiglu_limit: Optional[float] = None, + swiglu_alpha: Optional[float] = None, + swiglu_beta: Optional[float] = None, + ) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, + scale=scale, + dtype=dtype, + swiglu_limit=swiglu_limit, + swiglu_alpha=swiglu_alpha, + swiglu_beta=swiglu_beta, + ) + + def reference( + self, + x: torch.Tensor, + scale: Optional[torch.Tensor] = None, + dtype: Optional[torch.dtype] = None, + swiglu_limit: Optional[float] = None, + swiglu_alpha: Optional[float] = None, + swiglu_beta: Optional[float] = None, + ) -> torch.Tensor: + """fp32 throughout, one rounding to the output dtype; no matmul.""" + d = x.shape[-1] // 2 + gate = x[..., :d].float() + up = x[..., d:].float() + if swiglu_limit is not None and swiglu_limit > 0.0: + gate = gate.clamp(max=swiglu_limit) + up = up.clamp(-swiglu_limit, swiglu_limit) + alpha = 1.0 if swiglu_alpha is None else swiglu_alpha + beta = 0.0 if swiglu_beta is None else swiglu_beta + out = gate * torch.sigmoid(alpha * gate) * (up + beta) + out_dtype = x.dtype if dtype is None else dtype + if scale is not None: + bound = torch.finfo(out_dtype).max + out = (out / scale.reshape(-1)[0].float()).clamp(-bound, bound) + return out.to(out_dtype) + + def is_valid( + self, + x: torch.Tensor, + scale: Optional[torch.Tensor] = None, + dtype: Optional[torch.dtype] = None, + swiglu_limit: Optional[float] = None, + swiglu_alpha: Optional[float] = None, + swiglu_beta: Optional[float] = None, + ) -> None: + """Two inputs the op takes and answers wrongly, both read from the kernel source. + + A non-2-D `x` or an odd width the op rejects itself. + """ + # The kernel addresses column j of row i at `x + i * x.stride(0) + j`: the row stride + # is honoured, the column stride is assumed to be 1. + assert x.stride(-1) == 1, ( + f"x's last dim must be contiguous (stride 1), got stride {x.stride(-1)}; a strided " + "view is read as if dense and silently gives the wrong halves" + ) + # Without `dtype` the output stays in x's dtype, but the kernel still divides by the + # scale: scaled values land in an unquantized tensor. + assert scale is None or dtype is not None, ( + "scale without dtype divides the result by scale and stores it unquantized in " + "x's dtype; pass the fp8 dtype the scale quantizes to" + ) + + +silu_and_mul = SiluAndMul() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/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/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/compact_sparse_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/compact_sparse_rows.py new file mode 100644 index 000000000000..6bee98c4210a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/compact_sparse_rows.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pack each row's selected latent-cache row ids to the front, in one launch. + +The expanded selection of a sparse MLA row (``kpool_expand``) holds ``-1`` wherever a pool slot +or a tail slot selected nothing, while the native sparse attention reads each row's first +``lengths[r]`` ids. Per row this moves the nonnegative ids to the front in their order +(duplicates kept), pads with ``-1`` to the next multiple of 64 columns, and gives a row with no +valid id one dummy id ``0``, length 1 and ``nonempty`` False: the attention kernel takes no +empty row, so its caller discards that row's output. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +@triton.jit +def _compact_sparse_rows_kernel( + Indices, + Packed, + Lengths, + Nonempty, + Stride: tl.constexpr, + Width: tl.constexpr, + OutWidth: tl.constexpr, + Block: tl.constexpr, + OutBlock: tl.constexpr, +): + row = tl.program_id(0) + col = tl.arange(0, Block) + value = tl.load(Indices + row * Stride + col, col < Width, -1) + valid = (col < Width) & (value >= 0) + destination = tl.cumsum(valid.to(tl.int32)) - 1 + length = tl.sum(valid.to(tl.int32)) + tl.store(Packed + row * OutWidth + destination, value, valid) + tail = tl.arange(0, OutBlock) + tl.store( + Packed + row * OutWidth + tail, -1, (tail >= tl.maximum(length, 1)) & (tail < OutWidth) + ) + if length == 0: + tl.store(Packed + row * OutWidth, 0) + tl.store(Lengths + row, tl.maximum(length, 1)) + tl.store(Nonempty + row, length > 0) + + +class CompactSparseRows(OpWrapper): + """``indices`` int32 ``[R, W]`` -> (``packed`` int32 ``[R, cdiv(W, 64) * 64]``, ``lengths`` + int32 ``[R]``, ``nonempty`` bool ``[R]``). + + ``packed[r, :lengths[r]]`` are the ids ``>= 0`` of ``indices[r]`` in column order, the rest + of the row is ``-1``; a row with none is ``[0, -1, ...]`` with length 1 and ``nonempty`` + False. Negative ids are the only invalid ones: an id past the latent cache is kept. + """ + + ARCHS = frozenset({Arch.SM_100}) + # GLM-5.3-Flash selects index_topk = 2048 positions as 512 pools of 4 plus a tail of up to + # 3, so the expanded rows are kernel_output_width = cdiv(2051, 64) * 64 = 2112 wide. + # `kv` is the visible length of the cell's last row; the test builds each row's expansion + # from its own visible length. + CELLS: tuple[Cell, ...] = ( + Cell( + why="MTP-3 verify of one 13K-token request", + spec=dict(rows=4, width=2112, kv=13312), + ), + Cell( + why="MTP-3 verify of eight 13K-token requests", + spec=dict(rows=32, width=2112, kv=13312), + ), + Cell( + why="MTP-3 verify of 64 requests, the largest CUDA-graph batch", + spec=dict(rows=256, width=2112, kv=13312), + ), + Cell( + why="single-token decode of eight 13K-token requests", + spec=dict(rows=8, width=2112, kv=13312), + ), + Cell( + why="one rank's sequence-parallel quarter of a 13,312-token prompt (rows 0-3327)", + spec=dict(rows=3328, width=2112, kv=3328), + ), + Cell( + why="rank 1's quarter (rows 10-19) of a 37-token prompt: short rows", + spec=dict(rows=10, width=2112, kv=20), + ), + ) + + note = """ + The row stride of `indices` is a compile-time constant of the kernel, so every distinct + stride compiles once; the callers pass the contiguous output of `kpool_expand`. + + An id at or past the latent cache's row count is packed like any other: the native sparse + attention would read past the cache. `kpool_expand` writes in-range ids or -1 only. + """ + + def __call__(self, indices: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + rows = indices.shape[0] + width = triton.cdiv(indices.shape[1], 64) * 64 + packed = torch.empty((rows, width), dtype=torch.int32, device=indices.device) + lengths = torch.empty(rows, dtype=torch.int32, device=indices.device) + nonempty = torch.empty(rows, dtype=torch.bool, device=indices.device) + if rows == 0: + return packed, lengths, nonempty + _compact_sparse_rows_kernel[(rows,)]( + indices, + packed, + lengths, + nonempty, + indices.stride(0), + indices.shape[1], + width, + triton.next_power_of_2(indices.shape[1]), + triton.next_power_of_2(width), + ) + return packed, lengths, nonempty + + def reference(self, indices: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """A stable partition of each row (valid ids first, in order), then the padding and the + dummy id of an empty row.""" + rows, width = indices.shape + out_width = triton.cdiv(width, 64) * 64 + valid = indices >= 0 + # Stable sort on "is invalid" keeps the valid ids in their column order. + order = torch.sort((~valid).to(torch.int8), dim=1, stable=True).indices + moved = torch.gather(indices, 1, order) + count = valid.sum(dim=1) + col = torch.arange(out_width, device=indices.device) + packed = torch.full((rows, out_width), -1, dtype=torch.int32, device=indices.device) + packed[:, :width] = moved.to(torch.int32) + packed = torch.where(col[None, :] < count[:, None], packed, -1) + packed[:, 0] = torch.where(count > 0, packed[:, 0], 0) + return packed, count.clamp(min=1).to(torch.int32), count > 0 + + def is_valid(self, indices: torch.Tensor) -> None: + assert indices.dim() == 2 and indices.dtype == torch.int32, ( + f"indices must be int32 [R, W]; got {indices.dtype} {tuple(indices.shape)}: the " + "kernel stores the loaded ids into an int32 output" + ) + assert indices.stride(1) == 1, f"indices rows must be unit-stride; got {indices.stride()}" + # Row offsets are int32 products of the program id and a row width. + out_width = triton.cdiv(indices.shape[1], 64) * 64 + assert indices.shape[0] * max(indices.stride(0), out_width) < 2**31 + + +compact_sparse_rows = CompactSparseRows() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/head_major.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/head_major.py new file mode 100644 index 000000000000..0efad62909c6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/head_major.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""FP8 casts of head-major ``[H, R, D]`` batched-GEMM outputs into row-major ``[R, H * D]``. + +The MLA layers run their per-head products as ``torch.bmm`` over the heads, which returns +head-major tensors: the query into the latent space (generation rows with a rank's 16 heads, +sequence-parallel context rows with all 64) and the context rows' latent output into the value +space. Each consumer wants the rows contiguous and in FP8, so each cast reads the head-major +tensor once and writes the row-major FP8 copy, in place of a transposing copy followed by the +cast. + +``saturate_fp8``: the attention query, ``clamp(+-448)`` then E4M3 (the sparse attention takes an +E4M3 query). + +``quantize_1x128``: the context rows' o_proj input, the 1 x 128 block quantization +``fp8_quantize_1x128`` computes (scale ``448 / amax`` with ``amax`` floored at BF16 ``1e-10``, +FP32 product, saturating E4M3), with its K-block-major ``[H * D / 128, R]`` FP32 dequantization +scales. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..quantization.fp8_quantize_1x128 import fp8_quantize_1x128 + +_FP8_MAX = 448.0 + + +@triton.jit +def _saturate_kernel(x_ptr, out_ptr, R, H: tl.constexpr, D: tl.constexpr, BR: tl.constexpr): + h = tl.program_id(1) + r = tl.program_id(0) * BR + tl.arange(0, BR) + d = tl.arange(0, D) + mask = (r < R)[:, None] + x = tl.load(x_ptr + (h * R + r[:, None]).to(tl.int64) * D + d[None, :], mask=mask) + x = tl.minimum(tl.maximum(x.to(tl.float32), -448.0), 448.0) + tl.store( + out_ptr + (r[:, None].to(tl.int64) * H + h) * D + d[None, :], x.to(tl.float8e4nv), mask=mask + ) + + +@triton.jit +def _quantize_kernel( + x_ptr, out_ptr, scale_ptr, R, H: tl.constexpr, D: tl.constexpr, BR: tl.constexpr +): + h = tl.program_id(1) + r = tl.program_id(0) * BR + tl.arange(0, BR) + mask = (r < R)[:, None, None] + d = tl.arange(0, D).reshape(D // 128, 128) + x = tl.load(x_ptr + (h * R + r[:, None, None]).to(tl.int64) * D + d[None, :, :], mask=mask) + # The floor is 1e-10 rounded to BF16, as the quantizer compares in BF16. + amax = tl.maximum(tl.max(tl.abs(x), axis=2).to(tl.float32), 1.000444171950221e-10) + quant = tl.math.div_rn(tl.full(amax.shape, 448.0, tl.float32), amax) + q = (x.to(tl.float32) * quant[:, :, None]).to(tl.float8e4nv) + tl.store(out_ptr + (r[:, None, None].to(tl.int64) * H + h) * D + d[None, :, :], q, mask=mask) + blk = h * (D // 128) + tl.arange(0, D // 128) + dequant = tl.math.div_rn(tl.full(quant.shape, 1.0, tl.float32), quant) + tl.store(scale_ptr + blk[None, :] * R + r[:, None], dequant, mask=(r < R)[:, None]) + + +class SaturateFp8(OpWrapper): + """``clamp(x, -448, 448)`` as E4M3, ``[H, R, D]`` BF16 -> ``[R, H, D]``.""" + + ARCHS = frozenset({Arch.SM_100}) + CELLS: tuple[Cell, ...] = ( + Cell(why="MLA generation query, one draft row, 16 heads", spec=dict(h=16, r=1, d=512)), + Cell( + why="MLA generation query, MTP-3 verify of one request, 16 heads", + spec=dict(h=16, r=4, d=512), + ), + Cell( + why="MLA generation query, MTP-3 verify of eight requests, 16 heads", + spec=dict(h=16, r=32, d=512), + ), + Cell(why="MLA context query, 13K prompt / 4 ranks", spec=dict(h=64, r=3328, d=512)), + Cell(why="MLA context query, ragged rows", spec=dict(h=64, r=517, d=512)), + Cell(why="MLA context query, 37-token prompt / 4 ranks", spec=dict(h=64, r=10, d=512)), + ) + + def __call__(self, x: torch.Tensor) -> torch.Tensor: + h, r, d = x.shape + out = torch.empty((r, h, d), dtype=torch.float8_e4m3fn, device=x.device) + _saturate_kernel[(triton.cdiv(r, 8), h)](x, out, r, h, d, 8, num_warps=4) + return out + + def reference(self, x: torch.Tensor) -> torch.Tensor: + return x.transpose(0, 1).clamp(-_FP8_MAX, _FP8_MAX).to(torch.float8_e4m3fn) + + def is_valid(self, x: torch.Tensor) -> None: + assert x.dtype == torch.bfloat16 and x.is_contiguous() and x.dim() == 3 + assert x.shape[2] % 128 == 0 + + +class Quantize1x128(OpWrapper): + """``fp8_quantize_1x128`` of ``x.transpose(0, 1).reshape(R, H * D)`` for ``[H, R, D]`` BF16: + E4M3 ``[R, H * D]`` and FP32 ``[H * D / 128, R]``.""" + + ARCHS = frozenset({Arch.SM_100}) + CELLS: tuple[Cell, ...] = ( + Cell(why="MLA context values, 13K prompt / 4 ranks", spec=dict(h=64, r=3328, d=256)), + Cell(why="MLA context values, ragged rows", spec=dict(h=64, r=517, d=256)), + Cell(why="MLA context values, 37-token prompt / 4 ranks", spec=dict(h=64, r=10, d=256)), + ) + + def __call__(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + h, r, d = x.shape + out = torch.empty((r, h * d), dtype=torch.float8_e4m3fn, device=x.device) + scale = torch.empty((h * d // 128, r), dtype=torch.float32, device=x.device) + _quantize_kernel[(triton.cdiv(r, 8), h)](x, out, scale, r, h, d, 8, num_warps=4) + return out, scale + + def reference(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``fp8_quantize_1x128``'s plain-torch reference on the row-major copy.""" + h, r, d = x.shape + return fp8_quantize_1x128.reference(x.transpose(0, 1).reshape(r, h * d)) + + def is_valid(self, x: torch.Tensor) -> None: + assert x.dtype == torch.bfloat16 and x.is_contiguous() and x.dim() == 3 + assert x.shape[2] % 128 == 0 + + +saturate_fp8 = SaturateFp8() +quantize_1x128 = Quantize1x128() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/indexer_topk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/indexer_topk.py new file mode 100644 index 000000000000..d858b7b4f3fa --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/indexer_topk.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pool top-k of the GLM-5.3-Flash indexer: the CuTe DSL radix decode top-k with one candidate +count per row. + +Row ``r`` picks the ``select_k`` highest of its ``num_cand[r]`` complete-pool scores +``scores[r, :num_cand[r]]``. A row with ``num_cand[r] <= select_k`` keeps all its pools, +``0 .. num_cand[r] - 1`` in order and ``-1`` after them; a longer row gets ``select_k`` pool ids +in an unspecified order. Columns at and past ``num_cand[r]`` are not read. The target selects +the pools of its sequence-parallel context rows with it (``MlaLayer.context_rows``), on the +scores of ``kpool_score_fp8``; generation rows select with ``kpool_select``. The op is the CuTe +DSL radix decode top-k (``trtllm::cute_dsl_indexer_topk_decode``) with ``next_n = 1``. +""" + +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 _ordered_key(values: torch.Tensor) -> torch.Tensor: + """FP32 values as int64 keys in the order of their ordered bit patterns (-0 below +0).""" + bits = values.contiguous().view(torch.int32).long() + return torch.where(bits < 0, -(bits & 0x7FFFFFFF) - 1, bits) + + +class IndexerTopk(OpWrapper): + """``scores`` FP32 ``[R, C]``, ``num_cand`` int32 ``[R]`` -> ``selected`` int32 ``[R, + select_k]``, written in place and returned. + + ``selected`` is the caller's buffer, exactly ``select_k`` columns wide and contiguous. + ``scores`` needs unit-stride columns and may have padded rows (the FP8 scorer's strided + view); ``num_cand[r] <= C``. + """ + + ARCHS = frozenset({Arch.SM_100}) + # select_k = index_topk / index_kpool = 512. The target's context rows score the FP8 + # scorer's [rows, max_kv_len // 4] view, max_kv_len being the largest visible length of + # the scored rows, and DeepGEMM pads its rows (`row_pad` extra columns); context cells give + # a rank's sequence-parallel row range, whose candidate counts run up from (row + 1) // 4. + # The generation cells (`kv` the KV length after the step) drive the op at the 32,800-pool + # width of a full block table. `levels` rounds the scores to that many values per unit, so + # the k-th value is tied many times. + CELLS: tuple[Cell, ...] = ( + Cell( + why="MTP-3 verify of one 13K-token request, BF16-scorer width", + spec=dict(requests=1, tokens=4, kv=13312, width=32800), + ), + Cell( + why="MTP-3 verify of eight 13K-token requests, BF16-scorer width", + spec=dict(requests=8, tokens=4, kv=13312, width=32800), + ), + Cell( + why="MTP-3 verify of 64 requests, the largest CUDA-graph batch", + spec=dict(requests=64, tokens=4, kv=13312, width=32800), + ), + Cell( + why="single-token decode of eight 13K-token requests", + spec=dict(requests=8, tokens=1, kv=13312, width=32800), + ), + Cell( + why="rank 0's quarter of a 13,312-token prompt (rows 0-3327, 832 pools)", + spec=dict(rows=(0, 3328), width=832, row_pad=64), + ), + Cell( + why="last sequence-parallel quarter of a 13,312-token prompt (rows 9984-13311)", + spec=dict(rows=(9984, 13312), width=3328, row_pad=256), + ), + Cell( + why="rank 1's quarter (rows 10-19) of a 37-token prompt: five pools, all short rows", + spec=dict(rows=(10, 20), width=5, row_pad=3), + ), + Cell( + why="heavy ties at the k-th score", + spec=dict(requests=8, tokens=4, kv=13312, width=32800, levels=4), + ), + ) + + note = """ + The kernel's choice among pools of equal score at the k-th value is unspecified, and so + is the order of a long row's winners: the test compares each row's selected scores as a + sorted multiset, which every valid top-k shares, and the pool sets where no tie decides. + + JIT: the first call of a compile key compiles a CuTe DSL kernel for seconds. The key is + the score dtype, next_power_of_2(C), select_k and, through the dispatch on the row count, + the occupancy band (the cluster size or the large-occupancy variant). The target's C is + the FP8 context scorer's width, which follows the largest visible length of a scoring call, + so each power-of-two bucket of it, and each row band, compiles on its first prefill step + unless a warmup compiled it: `warmup_cute_dsl_radix_topk_decode(top_k=select_k, + num_cols=C, next_n=1)` in `tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops` compiles + every row band of one width. Registered only where the CuTe DSL is available. + + FP32 scores only: BF16 and FP16 scores take other kernels of the dispatch. + """ + + def __call__( + self, + scores: torch.Tensor, + num_cand: torch.Tensor, + selected: torch.Tensor, + *, + select_k: int, + ) -> torch.Tensor: + torch.ops.trtllm.cute_dsl_indexer_topk_decode(scores, num_cand, selected, select_k, 1) + return selected + + def reference( + self, + scores: torch.Tensor, + num_cand: torch.Tensor, + selected: torch.Tensor, + *, + select_k: int, + ) -> torch.Tensor: + """A fresh ``[R, select_k]`` int32 selection; ``selected`` is not read. + + A long row's winners are its ``select_k`` highest scores, ties at the k-th score broken + toward the lower pool id, listed in ascending pool order (the op's order is + unspecified). A short row is ``0 .. num_cand[r] - 1`` then ``-1``, as the op writes it. + """ + rows = scores.shape[0] + dev = scores.device + out = torch.full((rows, select_k), -1, dtype=torch.int32, device=dev) + count = num_cand.long() + ident = torch.arange(select_k, device=dev) + short = count <= select_k + out[short] = torch.where(ident < count[short, None], ident, -1).to(torch.int32) + for chunk in (~short).nonzero().flatten().split(64): + if chunk.numel() == 0: + continue + width = int(count[chunk].max()) + key = _ordered_key(scores[chunk, :width].float()) + col = torch.arange(width, device=dev) + key = torch.where(col < count[chunk, None], key, torch.iinfo(torch.int64).min) + # A stable descending sort puts the lower column first among equal keys. + order = torch.sort(key, dim=1, descending=True, stable=True).indices[:, :select_k] + out[chunk] = order.sort(dim=1).values.to(torch.int32) + return out + + def is_valid( + self, + scores: torch.Tensor, + num_cand: torch.Tensor, + selected: torch.Tensor, + *, + select_k: int, + ) -> None: + assert scores.dtype == torch.float32, f"FP32 scores only; got {scores.dtype}" + assert scores.dim() == 2 and scores.stride(1) == 1, f"scores strides {scores.stride()}" + # The compiled kernel assumes a 32-byte-aligned base for its vector loads. + assert scores.data_ptr() % 32 == 0, "scores must start 32-byte aligned" + rows = scores.shape[0] + # One int32 length per row: the kernel reads num_cand through an int32 pointer. + assert num_cand.dtype == torch.int32 and num_cand.shape == (rows,) + assert num_cand.is_contiguous() + # The kernel writes select_k columns per row at a row stride of select_k. + assert selected.dtype == torch.int32 and selected.shape == (rows, select_k) + assert selected.is_contiguous() + if rows: + # A count past the row reads the next row's scores. + assert 0 <= int(num_cand.min()) and int(num_cand.max()) <= scores.shape[1], ( + f"candidate counts up to {int(num_cand.max())} for {scores.shape[1]} columns" + ) + + +indexer_topk = IndexerTopk() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/kpool_expand.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/kpool_expand.py new file mode 100644 index 000000000000..5c1d24855eaa --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/kpool_expand.py @@ -0,0 +1,272 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Selected pools -> latent-cache row ids of their members and of the incomplete tail, in one +launch. + +Pool ``j`` of a row is its request's positions ``[j * kpool, (j + 1) * kpool)``; the row's +``kv_len // kpool`` complete pools are its candidates and the ``kv_len % kpool`` positions after +them (the incomplete tail) are always visible to it. Cache position ``p`` of block table ``t`` +is latent row ``base_row + block_tables[t, p // tokens_per_block] * rows_per_slot + p % +tokens_per_block`` of the uniformly strided row view of the latent pool. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +@triton.jit +def _kpool_expand_kernel( + SEL, + KV_LENS, + BT, + REQ, + OUT, + sel_stride, + bt_stride, + tpb, + base_row, + rows_per_slot, + out_width, + SELECT_K: tl.constexpr, + KPOOL: tl.constexpr, + KPOOL_P2: tl.constexpr, + PAD_BLOCK: tl.constexpr, + HAS_REQ: tl.constexpr, +): + """Expand selected pools into latent-cache row ids and append the tail. + + Output row ``b`` is ``[SELECT_K * KPOOL expanded members | KPOOL - 1 tail + positions | -1 padding]``; every invalid entry is ``-1`` (the FlashMLA + sparse kernel's own invalid marker). Row id of cache position ``p`` is + ``base_row + block_table[p // tpb] * rows_per_slot + p % tpb``. + """ + b = tl.program_id(0) + if HAS_REQ: + tbl = tl.load(REQ + b).to(tl.int64) + else: + tbl = b.to(tl.int64) + kv_len = tl.load(KV_LENS + b).to(tl.int64) + num_cand = kv_len // KPOOL + + i = tl.arange(0, SELECT_K) + j = tl.load(SEL + b * sel_stride + i).to(tl.int64) + # A radix top-k pads short rows with negative ids; those are invalid too. + valid = (j >= 0) & (j < num_cand) + m = tl.arange(0, KPOOL_P2) + mmask = m < KPOOL + pos = j[:, None] * KPOOL + m[None, :] + page = pos // tpb + off = pos - page * tpb + ok = valid[:, None] & mmask[None, :] + slot = tl.load(BT + tbl * bt_stride + page, mask=ok, other=0).to(tl.int64) + rows = base_row + slot * rows_per_slot + off + rows = tl.where(ok, rows, -1) + tl.store( + OUT + b * out_width + i[:, None] * KPOOL + m[None, :], + rows.to(tl.int32), + mask=mmask[None, :] & (i[:, None] >= 0), + ) + + # Tail: the incomplete trailing group, always visible to its own query. + tail_count = kv_len - num_cand * KPOOL + tail_start = kv_len - tail_count + t = tl.arange(0, KPOOL_P2) + tpos = tail_start + t + tvalid = t < tail_count + tpage = tpos // tpb + toff = tpos - tpage * tpb + tslot = tl.load(BT + tbl * bt_stride + tpage, mask=tvalid, other=0).to(tl.int64) + trows = tl.where(tvalid, base_row + tslot * rows_per_slot + toff, -1) + tidx = SELECT_K * KPOOL + t + tl.store(OUT + b * out_width + tidx, trows.to(tl.int32), mask=tidx < out_width) + + # -1 padding up to the kernel-aligned width. + pad0 = SELECT_K * KPOOL + KPOOL_P2 + for start in range(pad0, out_width, PAD_BLOCK): + p = start + tl.arange(0, PAD_BLOCK) + tl.store(OUT + b * out_width + p, tl.full([PAD_BLOCK], -1, tl.int32), mask=p < out_width) + + +class KpoolExpand(OpWrapper): + """``selected`` int32 ``[N, select_k]`` pool ids -> int32 ``[N, out_width]`` latent row ids. + + Column ``i * kpool + m`` holds member ``m`` of ``selected[n, i]`` when that id is a + candidate of row ``n`` (``0 <= id < kv_lens[n] // kpool``), else ``-1``; columns + ``select_k * kpool + t`` hold the tail positions ``(kv_lens[n] // kpool) * kpool + t`` for + ``t < kv_lens[n] % kpool``, then ``-1`` to ``out_width``. The order of ``selected`` is kept. + Row ``n`` reads block table ``request_ids[n]`` when given, else table ``n``. + """ + + ARCHS = frozenset({Arch.SM_100}) + # select_k = index_topk / index_kpool = 512 pools of kpool 4, out_width the padded + # kernel_output_width 2112, 32-token blocks. `kv` is the prompt length (context cells) or + # the requests' KV length after the step (generation cells). + CELLS: tuple[Cell, ...] = ( + Cell( + why="MTP-3 verify of one 13K-token request", + spec=dict(requests=1, tokens=4, kv=13312), + ), + Cell( + why="MTP-3 verify of eight 13K-token requests", + spec=dict(requests=8, tokens=4, kv=13312), + ), + Cell( + why="MTP-3 verify of 64 requests, the largest CUDA-graph batch", + spec=dict(requests=64, tokens=4, kv=13312), + ), + Cell( + why="single-token decode of eight requests: tables by row, no request ids", + spec=dict(requests=8, tokens=1, kv=13312), + ), + Cell( + why="one rank's sequence-parallel quarter of a 13,312-token prompt (rows 0-3327)", + spec=dict(context=((0, 13312),), rows=(0, 3328)), + ), + Cell( + why="a 13,312-token prompt's last sequence-parallel quarter (rows 9984-13311)", + spec=dict(context=((0, 13312),), rows=(9984, 13312)), + ), + Cell( + why="rank 1's quarter (rows 10-19) of a 37-token prompt: short rows", + spec=dict(context=((0, 37),), rows=(10, 20)), + ), + ) + + note = """ + The caller's `selected` comes from the decode top-k: candidate pool ids in an + unspecified order, padded with -1 when a row has fewer than select_k candidates. Ids at + or past a row's candidate count are written as -1 too. + + The kernel spans the selection with `tl.arange(0, select_k)`, so select_k is a power of + two (512 for this model); `out_width < select_k * kpool + kpool - 1` raises. + """ + + def __call__( + self, + selected: torch.Tensor, + kv_lens: torch.Tensor, + block_tables: torch.Tensor, + tokens_per_block: int, + *, + base_row: int, + rows_per_slot: int, + kpool: int, + out_width: int, + request_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + n, select_k = selected.shape + if out_width < select_k * kpool + kpool - 1: + raise ValueError( + f"kpool_expand: out_width {out_width} < {select_k * kpool + kpool - 1}" + ) + out = torch.empty(n, out_width, dtype=torch.int32, device=selected.device) + if n == 0: + return out + _kpool_expand_kernel[(n,)]( + selected, + kv_lens, + block_tables, + kv_lens if request_ids is None else request_ids, + out, + selected.stride(0), + block_tables.stride(0), + tokens_per_block, + base_row, + rows_per_slot, + out_width, + SELECT_K=select_k, + KPOOL=kpool, + KPOOL_P2=triton.next_power_of_2(kpool), + PAD_BLOCK=64, + HAS_REQ=request_ids is not None, + num_warps=4, + ) + return out + + def reference( + self, + selected: torch.Tensor, + kv_lens: torch.Tensor, + block_tables: torch.Tensor, + tokens_per_block: int, + *, + base_row: int, + rows_per_slot: int, + kpool: int, + out_width: int, + request_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + """The member and tail positions of each row, mapped through its block table.""" + n, select_k = selected.shape + dev = selected.device + out = torch.full((n, out_width), -1, dtype=torch.int32, device=dev) + if n == 0: + return out + table = torch.arange(n, device=dev) if request_ids is None else request_ids.long() + kv = kv_lens.long() + num_cand = kv // kpool + last_page = block_tables.shape[1] - 1 + + def row_ids(positions: torch.Tensor, valid: torch.Tensor) -> torch.Tensor: + page = (positions // tokens_per_block).clamp(0, last_page) + slot = block_tables[table.view(-1, *([1] * (positions.dim() - 1))), page].long() + ids = base_row + slot * rows_per_slot + positions % tokens_per_block + return torch.where(valid, ids, -1).to(torch.int32) + + pools = selected.long() + member = torch.arange(kpool, device=dev) + positions = pools[:, :, None] * kpool + member # [n, select_k, kpool] + valid = ((pools >= 0) & (pools < num_cand[:, None]))[:, :, None].expand_as(positions) + out[:, : select_k * kpool] = row_ids(positions.clamp(min=0), valid).view(n, -1) + tail = num_cand[:, None] * kpool + member # [n, kpool] + width = min(kpool, out_width - select_k * kpool) + tail_ids = row_ids(tail, tail < kv[:, None]) + out[:, select_k * kpool : select_k * kpool + width] = tail_ids[:, :width] + return out + + def is_valid( + self, + selected: torch.Tensor, + kv_lens: torch.Tensor, + block_tables: torch.Tensor, + tokens_per_block: int, + *, + base_row: int, + rows_per_slot: int, + kpool: int, + out_width: int, + request_ids: torch.Tensor | None = None, + ) -> None: + n, select_k = selected.shape + assert selected.dtype == torch.int32 and selected.stride(1) == 1, ( + f"selected must be int32 with unit-stride rows; got {selected.dtype} " + f"{selected.stride()}" + ) + for name, width in (("select_k", select_k), ("kpool", kpool)): + assert width > 0 and width & (width - 1) == 0, f"{name} {width} must be a power of 2" + assert kv_lens.shape == (n,) + assert block_tables.dim() == 2 and block_tables.stride(1) == 1 + if n == 0: + return + if request_ids is None: + assert block_tables.shape[0] >= n, "row n reads block table n" + else: + assert request_ids.shape == (n,) + assert 0 <= int(request_ids.min()) and int(request_ids.max()) < block_tables.shape[0] + # A visible position past the table's width reads the next table's pages. + widest = int(kv_lens.max()) + assert widest <= block_tables.shape[1] * tokens_per_block, ( + f"a visible length of {widest} exceeds the tables' " + f"{block_tables.shape[1] * tokens_per_block} positions" + ) + # Row ids are int32. + top = base_row + (int(block_tables.max()) + 1) * rows_per_slot + assert top < 2**31, f"latent row ids up to {top} overflow int32" + + +kpool_expand = KpoolExpand() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/kpool_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/kpool_score.py new file mode 100644 index 000000000000..ef6fe76ee247 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/kpool_score.py @@ -0,0 +1,265 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pool scores of the GLM-5.3-Flash indexer over the paged pool keys, BF16 tensor-core dots, in +one launch. + +Row ``r`` scores each of its ``kv_lens[r] // kpool`` complete pools ``j``: + + score[r, j] = sum_h (w[r, h] * w_scale) * relu(q_scale * q[r, h, :] . key_j) + +where ``key_j`` is the pool key cached in the last ``head_dim`` columns of the indexer row of the +pool's first member (``kpool_update`` maintains it). BF16 x BF16 products are exact in FP32, so +each dot is an FP32 accumulation of exact products. Pools a row cannot see yet get the FP32 +minimum, so a following top-k never prefers them. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +#: The score of a pool a row cannot see. +FP32_MIN = torch.finfo(torch.float32).min + + +@triton.jit +def _kpool_score_kernel( + Q, + W, + POOL, + BT, + KV_LENS, + OUT, + w_row_stride, + w_head_stride, + slot_stride, + row_stride, + bt_stride, + tpb, + num_pools_max, + q_scale, + w_scale, + min_value, + HD: tl.constexpr, + KPOOL: tl.constexpr, + H: tl.constexpr, + HP: tl.constexpr, + BP: tl.constexpr, +): + """Score ``BP`` pools of row ``r`` against block table ``r``: ``sum_h w[h] * relu(q_h . key_j + * s)``, BF16 operands into the tensor-core dot with FP32 accumulation. + + Pools whose last member is not yet visible to the row (``j >= kv_len // KPOOL``) get + ``min_value`` so a following top-k never picks them ahead of a candidate; a block of pools + past every visible one stores ``min_value`` without gathering keys. + """ + r = tl.program_id(0).to(tl.int64) + j0 = tl.program_id(1) * BP + j = j0 + tl.arange(0, BP) + out_row = OUT + r * num_pools_max + num_cand = tl.load(KV_LENS + r).to(tl.int64) // KPOOL + if j0 >= num_cand: + tl.store(out_row + j, tl.full([BP], min_value, tl.float32), mask=j < num_pools_max) + else: + d = tl.arange(0, HD) + h = tl.arange(0, HP) + hmask = h < H + valid = j < num_cand + pos = j.to(tl.int64) * KPOOL + page = pos // tpb + off = pos - page * tpb + slot = tl.load(BT + r * bt_stride + page, mask=valid, other=0).to(tl.int64) + kptr = POOL + slot * slot_stride + off * row_stride + 2 * HD + keys = tl.load(kptr[:, None] + d[None, :], mask=valid[:, None], other=0.0).to(tl.bfloat16) + q = tl.load( + Q + r * H * HD + h[:, None] * HD + d[None, :], mask=hmask[:, None], other=0.0 + ).to(tl.bfloat16) + scores = tl.dot(q, tl.trans(keys), input_precision="ieee") # [HP, BP] + scores = tl.maximum(scores * q_scale, 0.0) + w = tl.load(W + r * w_row_stride + h * w_head_stride, mask=hmask, other=0.0) + mixed = tl.sum(scores * (w.to(tl.float32) * w_scale)[:, None], axis=0) + tl.store(out_row + j, tl.where(valid, mixed, min_value), mask=j < num_pools_max) + + +class KpoolScore(OpWrapper): + """``q`` BF16 ``[N, H, head_dim]`` (contiguous), ``weights`` ``[N, H]`` (any strides) -> + FP32 ``[N, num_pools_max]`` pool scores. + + Row ``r`` reads block table ``r`` and sees ``kv_lens[r]`` positions: columns ``[0, + kv_lens[r] // kpool)`` are its scores, every other column is the FP32 minimum. + """ + + ARCHS = frozenset({Arch.SM_100}) + # The indexer has 32 heads of 128 and scores pools of 4 over 32-token blocks. The tables + # are 4,100 pages wide, so the capacity is num_pools_max = 32,800 pools. The cells give the + # requests' KV length after the step; a verify step scores its rows on per-row copies of + # the generation tables (`GenerationRows.tables`). + CELLS: tuple[Cell, ...] = ( + Cell( + why="MTP-3 verify of one 13K-token request", + spec=dict(requests=1, tokens=4, kv=13312), + ), + Cell( + why="MTP-3 verify of eight 13K-token requests", + spec=dict(requests=8, tokens=4, kv=13312), + ), + Cell( + why="MTP-3 verify of 64 requests, the largest CUDA-graph batch", + spec=dict(requests=64, tokens=4, kv=13312), + ), + Cell( + why="single-token decode of eight 13K-token requests", + spec=dict(requests=8, tokens=1, kv=13312), + ), + Cell( + why="one row: an MTP draft step of one 13K-token request", + spec=dict(requests=1, tokens=1, kv=13312), + ), + ) + + note = """ + The target scores its generation rows here (`MlaLayer.generation_partial`); packed context + rows are scored by `kpool_score_fp8`. + + Only BF16 queries and BF16 pool keys: the dot reads both as BF16, so an FP32 query would + be rounded to BF16 silently. `is_valid` refuses it. + """ + + def __call__( + self, + q: torch.Tensor, + weights: torch.Tensor, + index_pool: torch.Tensor, + block_tables: torch.Tensor, + kv_lens: torch.Tensor, + tokens_per_block: int, + *, + num_pools_max: int, + head_dim: int, + kpool: int, + q_scale: float, + w_scale: float, + ) -> torch.Tensor: + n, num_heads, hd = q.shape + if hd != head_dim: + raise ValueError(f"kpool_score: q head_dim {hd} != {head_dim}") + out = torch.empty(n, num_pools_max, dtype=torch.float32, device=q.device) + if n == 0: + return out + bp = 64 + _kpool_score_kernel[(n, triton.cdiv(num_pools_max, bp))]( + q, + weights, + index_pool, + block_tables, + kv_lens, + out, + weights.stride(0), + weights.stride(1), + index_pool.stride(0), + index_pool.stride(1), + block_tables.stride(0), + tokens_per_block, + num_pools_max, + q_scale, + w_scale, + FP32_MIN, + HD=head_dim, + KPOOL=kpool, + H=num_heads, + HP=max(16, triton.next_power_of_2(num_heads)), + BP=bp, + num_warps=4, + ) + return out + + def reference( + self, + q: torch.Tensor, + weights: torch.Tensor, + index_pool: torch.Tensor, + block_tables: torch.Tensor, + kv_lens: torch.Tensor, + tokens_per_block: int, + *, + num_pools_max: int, + head_dim: int, + kpool: int, + q_scale: float, + w_scale: float, + ) -> torch.Tensor: + """The scores in FP32 from the BF16 operands, every invisible column the FP32 minimum.""" + n = q.shape[0] + dev = q.device + out = torch.full((n, num_pools_max), FP32_MIN, dtype=torch.float32, device=dev) + if n == 0: + return out + table = torch.arange(n, device=dev) + pools = (kv_lens.long() // kpool).clamp(max=num_pools_max) + pool_keys = index_pool[:, :, 2 * head_dim : 3 * head_dim] + chunk = 32 + for r0 in range(0, n, chunk): + rows = slice(r0, min(r0 + chunk, n)) + width = int(pools[rows].max()) + if width == 0: + continue + first = torch.arange(width, device=dev) * kpool + slot = block_tables[table[rows, None], (first // tokens_per_block)[None, :]] + keys = pool_keys[slot.long(), (first % tokens_per_block)[None, :]].float() + with true_fp32_matmul(): + dots = torch.einsum("rhd,rpd->rhp", q[rows].float(), keys) + mixed = (dots * q_scale).relu() * (weights[rows].float() * w_scale)[:, :, None] + scores = mixed.sum(dim=1) + col = torch.arange(width, device=dev) + out[rows, :width] = torch.where(col < pools[rows, None], scores, FP32_MIN) + return out + + def is_valid( + self, + q: torch.Tensor, + weights: torch.Tensor, + index_pool: torch.Tensor, + block_tables: torch.Tensor, + kv_lens: torch.Tensor, + tokens_per_block: int, + *, + num_pools_max: int, + head_dim: int, + kpool: int, + q_scale: float, + w_scale: float, + ) -> None: + assert q.dtype == torch.bfloat16 and index_pool.dtype == torch.bfloat16, ( + f"the dot reads q and the pool keys as BF16; got {q.dtype}, {index_pool.dtype}" + ) + # The kernel reads q at r * H * head_dim + h * head_dim + d. + assert q.dim() == 3 and q.is_contiguous(), f"q must be a contiguous [N, H, D]; {q.stride()}" + n, num_heads, _ = q.shape + assert head_dim & (head_dim - 1) == 0, f"head_dim {head_dim} must be a power of two" + assert weights.shape == (n, num_heads) + assert index_pool.dim() == 3 and index_pool.stride(2) == 1 + assert index_pool.shape[1] == tokens_per_block and index_pool.shape[2] >= 3 * head_dim + assert block_tables.dim() == 2 and block_tables.stride(1) == 1 + assert kv_lens.shape == (n,) + if n == 0: + return + # Row r reads table r. + assert block_tables.stride(0) == 0 or block_tables.shape[0] >= n + widest = int(kv_lens.max()) + # A row with more pools than num_pools_max loses the rest without a sign. + assert widest // kpool <= num_pools_max, f"{widest // kpool} pools > {num_pools_max}" + assert widest <= block_tables.shape[1] * tokens_per_block, ( + f"a visible length of {widest} exceeds the tables' " + f"{block_tables.shape[1] * tokens_per_block} positions" + ) + + +kpool_score = KpoolScore() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/kpool_score_fp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/kpool_score_fp8.py new file mode 100644 index 000000000000..e0325fe0be38 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/kpool_score_fp8.py @@ -0,0 +1,229 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pool scores of packed context rows in FP8, through DeepGEMM's MQA logits. + +The same scores as ``kpool_score``, + + score[r, j] = sum_h (w[r, h] * w_scale) * relu(q_scale * q[r, h, :] . key_j), + +for the context rows of one prefill step, computed on E4M3 operands: the pool keys of every +block table are gathered into ``[tables * pools, head_dim]`` and quantized per key, the queries +per (row, head), each with an FP32 scale ``amax / 448`` (``fused_cat_fp8``). The query scales +and ``w_scale * q_scale`` fold into FP32 head weights, and ``fp8_mqa_logits`` scores row ``r`` +against the key range of its table ``request_ids[r]``. The launch sequence is the gather, two +quantizations, the head-weight product, the logits and, for several tables, a per-row copy. +""" + +from __future__ import annotations + +import torch +import triton + +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.deep_gemm import fp8_mqa_logits + +#: The score the reference gives a pool a row cannot see; the op leaves those columns unspecified. +FP32_MIN = torch.finfo(torch.float32).min + + +class KpoolScoreFp8(OpWrapper): + """``q`` BF16 ``[N, H, 128]``, ``weights`` ``[N, H]`` -> FP32 scores, or None. + + Row ``r`` scores the first ``kv_lens[r] // kpool`` pools of table ``request_ids[r]``; its + later columns are not meaningful (unwritten or another table's keys), so a reader stops at + ``kv_lens[r] // kpool``. ``max_kv_len`` is a host bound on every ``kv_lens[r]`` and sets the + ``pools = max_kv_len // kpool`` gathered per table. The result is + + - one table: ``[N, pools]``, a strided view of DeepGEMM's logits whose row stride is the + kernel's aligned row width (a multiple of 4 columns, wider than ``pools``); + - several tables: a fresh ``[N, max(1024, cdiv(pools, 64) * 64)]`` tensor with each row's + ``pools`` columns copied in front, the rest unwritten; + - None when there is no row or no complete pool (``max_kv_len < kpool``). + """ + + ARCHS = frozenset({Arch.SM_100}) + # The indexer has 32 heads of 128 and scores pools of 4 over 32-token blocks. Each cell + # gives (cached, new) per context request and the scored row range (a rank's ceil(T / 4) + # sequence-parallel rows); max_kv_len is the largest visible length among those rows, as + # `MlaLayer.context_rows` passes it. The two-request cell takes the copy path and puts the + # second request's rows at another table's key range. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one rank's sequence-parallel quarter of a 13,312-token prompt (rows 0-3327)", + spec=dict(context=((0, 13312),), rows=(0, 3328)), + ), + Cell( + why="a 13,312-token prompt's last sequence-parallel quarter (rows 9984-13311)", + spec=dict(context=((0, 13312),), rows=(9984, 13312)), + ), + Cell( + why="rank 1's quarter (rows 10-19) of a 37-token prompt: two to five pools", + spec=dict(context=((0, 37),), rows=(10, 20)), + ), + Cell( + why="two packed context requests, a chunk continuation and a fresh prompt", + spec=dict(context=((4099, 3000), (0, 1313)), rows=(0, 4313)), + ), + ) + + note = """ + The scores are those of the E4M3-rounded operands. With per-row scales `amax / 448`, an + operand's relative rounding error is at most 2**-4 and a score moves by a few percent of + its scale; a score whose operands are on the E4M3 grid is exact up to FP32 accumulation. + The pools selected from these scores can therefore differ from those of the BF16 scores + (`kpool_score`) where two pools score within that rounding of each other. + + sm_100 only (DeepGEMM MQA logits). The target scores every packed context row here + (`MlaLayer.context_rows`); when this returns None the rows have no complete pool and + attend their incomplete trailing pool only. + + Transients grow with tables * pools: the gathered BF16 keys, their FP8 copy and scale, + the padded logits and, for several tables, the row copy. The target splits packed context + rows into calls of bounded size (`_score_calls` in `layers/mla.py`); this entry takes + whatever it is given. + """ + + def __call__( + self, + q: torch.Tensor, + weights: torch.Tensor, + index_pool: torch.Tensor, + block_tables: torch.Tensor, + kv_lens: torch.Tensor, + tokens_per_block: int, + *, + head_dim: int, + kpool: int, + q_scale: float, + w_scale: float, + request_ids: torch.Tensor, + max_kv_len: int, + ) -> torch.Tensor | None: + n = q.shape[0] + pools = max(0, int(max_kv_len)) // kpool + tables = block_tables.shape[0] + if n == 0 or pools == 0: + return None + positions = torch.arange(pools, device=q.device) * kpool + page = positions // tokens_per_block + offset = positions - page * tokens_per_block + slots = block_tables[:, page].long() # [tables, pools] + keys = index_pool[slots, offset, 2 * head_dim : 3 * head_dim].reshape(-1, head_dim) + half = head_dim // 2 + keys_fp8, key_scale = torch.ops.trtllm.fused_cat_fp8(keys[:, :half], keys[:, half:], False) + q_rows = q.reshape(-1, head_dim) + q_fp8, q_row_scale = torch.ops.trtllm.fused_cat_fp8( + q_rows[:, :half], q_rows[:, half:], False + ) + head_weights = weights.float() * q_row_scale.view(n, -1) * (w_scale * q_scale) + start = request_ids.to(torch.int32) * pools + end = start + (kv_lens // kpool).to(torch.int32) + logits = fp8_mqa_logits( + q_fp8.view(q.shape), + (keys_fp8, key_scale.view(-1)), + head_weights, + start, + end, + clean_logits=False, + ) + if tables == 1: + # DeepGEMM pads the logits rows to an aligned stride, as the pool top-k needs. + return logits[:, :pools] + # Keep the copy's rows aligned and wider than the selection, like the logits arena. + out = torch.empty( + n, max(1024, triton.cdiv(pools, 64) * 64), dtype=torch.float32, device=q.device + ) + rows = torch.arange(n, device=q.device) + out[:, :pools] = logits[:, : tables * pools].view(n, tables, pools)[ + rows, request_ids.long() + ] + return out + + def reference( + self, + q: torch.Tensor, + weights: torch.Tensor, + index_pool: torch.Tensor, + block_tables: torch.Tensor, + kv_lens: torch.Tensor, + tokens_per_block: int, + *, + head_dim: int, + kpool: int, + q_scale: float, + w_scale: float, + request_ids: torch.Tensor, + max_kv_len: int, + ) -> torch.Tensor | None: + """The exact scores the FP8 call approximates, FP32 from the BF16 operands, as a + contiguous ``[N, pools]``: FP32 minimum past each row's ``kv_lens[r] // kpool``, None + where the op returns None.""" + n = q.shape[0] + pools = max(0, int(max_kv_len)) // kpool + if n == 0 or pools == 0: + return None + dev = q.device + out = torch.full((n, pools), FP32_MIN, dtype=torch.float32, device=dev) + first = torch.arange(pools, device=dev) * kpool + slot = block_tables[:, first // tokens_per_block].long() # [tables, pools] + keys = index_pool[:, :, 2 * head_dim : 3 * head_dim][slot, (first % tokens_per_block)] + keys = keys.float() # [tables, pools, head_dim] + table = request_ids.long() + visible = kv_lens.long() // kpool + col = torch.arange(pools, device=dev) + for r0 in range(0, n, 64): + rows = slice(r0, min(r0 + 64, n)) + with true_fp32_matmul(): + dots = torch.einsum("rhd,rpd->rhp", q[rows].float(), keys[table[rows]]) + mixed = (dots * q_scale).relu() * (weights[rows].float() * w_scale)[:, :, None] + scores = mixed.sum(dim=1) + out[rows] = torch.where(col < visible[rows, None], scores, FP32_MIN) + return out + + def is_valid( + self, + q: torch.Tensor, + weights: torch.Tensor, + index_pool: torch.Tensor, + block_tables: torch.Tensor, + kv_lens: torch.Tensor, + tokens_per_block: int, + *, + head_dim: int, + kpool: int, + q_scale: float, + w_scale: float, + request_ids: torch.Tensor, + max_kv_len: int, + ) -> None: + assert q.dim() == 3 and q.shape[2] == head_dim + n, num_heads, _ = q.shape + assert weights.shape == (n, num_heads) + assert index_pool.dim() == 3 and index_pool.shape[1] == tokens_per_block + assert index_pool.shape[2] >= 3 * head_dim + assert block_tables.dim() == 2 + assert kv_lens.shape == (n,) and request_ids.shape == (n,) + pools = max(0, int(max_kv_len)) // kpool + if n == 0 or pools == 0: + return + # Row r's key range is [request_ids[r] * pools, + kv_lens[r] // kpool): a visible + # length past max_kv_len, or a table id past the tables, reads another table's keys. + widest = int(kv_lens.max()) + assert widest <= max_kv_len, f"kv_lens reach {widest} > max_kv_len {max_kv_len}" + assert int(kv_lens.min()) >= 0 + assert 0 <= int(request_ids.min()) and int(request_ids.max()) < block_tables.shape[0], ( + f"request_ids must index the {block_tables.shape[0]} tables passed" + ) + last_page = (pools - 1) * kpool // tokens_per_block + assert last_page < block_tables.shape[1], ( + f"max_kv_len {max_kv_len} reaches page {last_page} of {block_tables.shape[1]}" + ) + + +kpool_score_fp8 = KpoolScoreFp8() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/kpool_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/kpool_select.py new file mode 100644 index 000000000000..9c3a73577386 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/kpool_select.py @@ -0,0 +1,294 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pool top-k of a sparse MLA decode step, expanded to packed latent-cache rows, in one launch. + +Per row: the ``select_k`` highest pool scores among the row's ``kv_len // kpool`` complete +pools (an exact radix select on the FP32 bits), each selected pool expanded to its ``kpool`` +latent-cache row ids, then the rows of the incomplete trailing pool, packed from column 0 in +pool order and padded with -1. A row with no visible key gets one dummy id 0, length 1 and +``nonempty`` False. The result is that of ``kpool_expand`` and ``compact_sparse_rows`` on the +pools a top-k selects (up to the choice among pools of equal score at the threshold), with the +selected pools in ascending order; the target's context rows take those three steps +(``indexer_topk``, ``kpool_expand``, ``compact_sparse_rows``), its generation rows this one +launch (``MlaLayer.generation_partial``). +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .compact_sparse_rows import compact_sparse_rows +from .kpool_expand import kpool_expand + + +@triton.jit +def _order_key(scores): + """FP32 -> uint32 with the same order (larger float, larger key).""" + bits = scores.to(tl.uint32, bitcast=True) + return tl.where((bits >> 31) != 0, ~bits, bits | 0x80000000) + + +@triton.jit +def _in_prefix(key, thresh, shift): + """Whether ``key`` agrees with ``thresh`` on the bits above ``shift + 8`` (all of them in the + first pass; two shifts keep every shift amount below 32).""" + return ((key >> (shift + 7)) >> 1) == ((thresh >> (shift + 7)) >> 1) + + +@triton.jit +def _radix_step(hist, remaining, thresh, shift): + """Fix the 8 key bits at ``shift`` of the ``remaining``-th largest key from their histogram.""" + bins = tl.arange(0, 256) + at_or_above = tl.cumsum(hist, 0, reverse=True) + b = tl.max(tl.where(at_or_above >= remaining, bins, -1), 0) + remaining -= tl.sum(tl.where(bins > b, hist, 0), 0) + return remaining, thresh | (b.to(tl.uint32) << shift) + + +@triton.jit +def _select_in_registers(srow, sel_row, num_cand, SELECT_K: tl.constexpr, WIDTH: tl.constexpr): + """Rows with at most ``WIDTH`` candidates: one load, the radix passes and the emit on the + register tile. Writes the selected pool ids ascending to ``sel_row``; returns their count.""" + idx = tl.arange(0, WIDTH) + m = idx < num_cand + key = _order_key(tl.load(srow + idx, mask=m, other=0.0)) + thresh = tl.zeros([], tl.uint32) + ties = tl.full([], 1 << 30, tl.int32) + if num_cand > SELECT_K: + remaining = tl.full([], SELECT_K, tl.int32) + # Runtime pass loop: an unrolled one quadruples the code, which is fetched cold on + # every decode layer. + for p in range(4): + shift = (24 - 8 * p).to(tl.uint32) + match = m & _in_prefix(key, thresh, shift) + hist = tl.histogram(((key >> shift) & 255).to(tl.int32), 256, mask=match) + remaining, thresh = _radix_step(hist, remaining, thresh, shift) + ties = remaining + eq = m & (key == thresh) + sel = (m & (key > thresh)) | (eq & (tl.cumsum(eq.to(tl.int32), 0) <= ties)) + dest = tl.cumsum(sel.to(tl.int32), 0) - 1 + tl.store(sel_row + dest, idx.to(tl.int32), mask=sel) + return tl.sum(sel.to(tl.int32), 0) + + +@triton.jit +def _select_chunked(srow, sel_row, num_cand, SELECT_K: tl.constexpr, CHUNK: tl.constexpr): + """Longer rows: every radix pass and the emit stream the scores in ``CHUNK``-wide pieces.""" + lane = tl.arange(0, CHUNK) + thresh = tl.zeros([], tl.uint32) + remaining = tl.full([], SELECT_K, tl.int32) + for p in range(4): + shift = (24 - 8 * p).to(tl.uint32) + hist = tl.zeros([256], tl.int32) + for c0 in range(0, num_cand, CHUNK): + idx = c0 + lane + m = idx < num_cand + key = _order_key(tl.load(srow + idx, mask=m, other=0.0)) + m = m & _in_prefix(key, thresh, shift) + hist += tl.histogram(((key >> shift) & 255).to(tl.int32), 256, mask=m) + remaining, thresh = _radix_step(hist, remaining, thresh, shift) + n_sel = tl.zeros([], tl.int32) + ties_seen = tl.zeros([], tl.int32) + for c0 in range(0, num_cand, CHUNK): + idx = c0 + lane + m = idx < num_cand + key = _order_key(tl.load(srow + idx, mask=m, other=0.0)) + eq = m & (key == thresh) + eq_rank = ties_seen + tl.cumsum(eq.to(tl.int32), 0) + sel = (m & (key > thresh)) | (eq & (eq_rank <= remaining)) + dest = n_sel + tl.cumsum(sel.to(tl.int32), 0) - 1 + tl.store(sel_row + dest, idx.to(tl.int32), mask=sel) + ties_seen += tl.sum(eq.to(tl.int32), 0) + n_sel += tl.sum(sel.to(tl.int32), 0) + return n_sel + + +@triton.jit +def _kpool_select_kernel( + SCORES, + KV_LENS, + BT, + SEL, + PACKED, + LENGTHS, + NONEMPTY, + score_stride, + bt_stride, + tpb, + base_row, + rows_per_slot, + SELECT_K: tl.constexpr, + KPOOL: tl.constexpr, + KPOOL_P2: tl.constexpr, + OUT_W: tl.constexpr, + PAD_BLOCK: tl.constexpr, +): + r = tl.program_id(0).to(tl.int64) + kv_len = tl.load(KV_LENS + r).to(tl.int64) + num_cand = kv_len // KPOOL + srow = SCORES + r * score_stride + sel_row = SEL + r * SELECT_K + # Selected pool ids, ascending, into SEL[r, :n_sel]: on one register tile up to 4096 + # candidates (a 16K-token context at kpool 4), streamed in chunks past that. + if num_cand <= 4096: + n_sel = _select_in_registers(srow, sel_row, num_cand, SELECT_K, 4096) + else: + n_sel = _select_chunked(srow, sel_row, num_cand, SELECT_K, 4096) + tl.debug_barrier() + + # Expand: selected pool i -> columns [i * KPOOL, (i + 1) * KPOOL). + i = tl.arange(0, SELECT_K) + mm = tl.arange(0, KPOOL_P2) + ok = (i < n_sel)[:, None] & (mm < KPOOL)[None, :] + j = tl.load(SEL + r * SELECT_K + i, mask=i < n_sel, other=0).to(tl.int64) + pos = j[:, None] * KPOOL + mm[None, :] + page = pos // tpb + slot = tl.load(BT + r * bt_stride + page, mask=ok, other=0).to(tl.int64) + rows = base_row + slot * rows_per_slot + (pos - page * tpb) + tl.store(PACKED + r * OUT_W + i[:, None] * KPOOL + mm[None, :], rows.to(tl.int32), mask=ok) + + # The incomplete trailing pool, always visible to its own query. + tail_count = kv_len - num_cand * KPOOL + t = tl.arange(0, KPOOL_P2) + tvalid = t < tail_count + tpos = num_cand * KPOOL + t + tpage = tpos // tpb + tslot = tl.load(BT + r * bt_stride + tpage, mask=tvalid, other=0).to(tl.int64) + trows = base_row + tslot * rows_per_slot + (tpos - tpage * tpb) + length = n_sel * KPOOL + tail_count.to(tl.int32) + tl.store(PACKED + r * OUT_W + n_sel * KPOOL + t, trows.to(tl.int32), mask=tvalid) + + # -1 after the packed ids; a row with none gets the dummy id 0. + start = tl.maximum(length, 1) + for c0 in range(0, OUT_W, PAD_BLOCK): + col = c0 + tl.arange(0, PAD_BLOCK) + tl.store( + PACKED + r * OUT_W + col, + tl.full([PAD_BLOCK], -1, tl.int32), + mask=(col >= start) & (col < OUT_W), + ) + if length == 0: + tl.store(PACKED + r * OUT_W, 0) + tl.store(LENGTHS + r, tl.maximum(length, 1)) + tl.store(NONEMPTY + r, length > 0) + + +class KpoolSelect(OpWrapper): + """``scores`` FP32 [R, >= max candidates] (``kpool_score``), ``kv_lens`` [R], ``block_tables`` + [R, blocks] -> (``packed`` int32 [R, W], ``lengths`` int32 [R], ``nonempty`` bool [R]) with + ``W = cdiv(out_width, 64) * 64``, the layout ``compact_sparse_rows`` returns.""" + + ARCHS = frozenset({Arch.SM_100}) + CELLS: tuple[Cell, ...] = ( + Cell(why="MTP-3 verify of one 13K-token request", spec=dict(rows=4, kv=13312, pools=32800)), + Cell( + why="MTP-3 verify of eight 13K-token requests", + spec=dict(rows=32, kv=13312, pools=32800), + ), + Cell( + why="short contexts: fewer complete pools than select_k", + spec=dict(rows=4, kv=1500, pools=32800), + ), + Cell( + why="20K-token context: chunked radix passes", spec=dict(rows=4, kv=20000, pools=32800) + ), + Cell(why="128K-token context", spec=dict(rows=2, kv=131072, pools=32800)), + ) + + def __call__( + self, + scores: torch.Tensor, + kv_lens: torch.Tensor, + block_tables: torch.Tensor, + tokens_per_block: int, + *, + select_k: int, + kpool: int, + base_row: int, + rows_per_slot: int, + out_width: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + rows = scores.shape[0] + width = triton.cdiv(out_width, 64) * 64 + packed = torch.empty((rows, width), dtype=torch.int32, device=scores.device) + lengths = torch.empty(rows, dtype=torch.int32, device=scores.device) + nonempty = torch.empty(rows, dtype=torch.bool, device=scores.device) + selected = torch.empty((rows, select_k), dtype=torch.int32, device=scores.device) + _kpool_select_kernel[(rows,)]( + scores, + kv_lens, + block_tables, + selected, + packed, + lengths, + nonempty, + scores.stride(0), + block_tables.stride(0), + tokens_per_block, + base_row, + rows_per_slot, + SELECT_K=select_k, + KPOOL=kpool, + KPOOL_P2=triton.next_power_of_2(kpool), + OUT_W=width, + PAD_BLOCK=512, + num_warps=16, + ) + return packed, lengths, nonempty + + def reference( + self, + scores: torch.Tensor, + kv_lens: torch.Tensor, + block_tables: torch.Tensor, + tokens_per_block: int, + *, + select_k: int, + kpool: int, + base_row: int, + rows_per_slot: int, + out_width: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + num_cand = (kv_lens.long() // kpool).clamp(max=scores.shape[1]) + col = torch.arange(scores.shape[1], device=scores.device) + masked = torch.where(col[None, :] < num_cand[:, None], scores, float("-inf")) + top = torch.topk(masked, select_k, dim=1).indices + top = torch.where(torch.gather(masked, 1, top) > float("-inf"), top, -1) + top = torch.sort(torch.where(top >= 0, top, scores.shape[1]), dim=1).values + top = torch.where(top < scores.shape[1], top, -1).int() + expanded = kpool_expand.reference( + top, + kv_lens, + block_tables, + tokens_per_block, + base_row=base_row, + rows_per_slot=rows_per_slot, + kpool=kpool, + out_width=out_width, + ) + return compact_sparse_rows.reference(expanded) + + def is_valid( + self, + scores: torch.Tensor, + kv_lens: torch.Tensor, + block_tables: torch.Tensor, + tokens_per_block: int, + *, + select_k: int, + kpool: int, + base_row: int, + rows_per_slot: int, + out_width: int, + ) -> None: + assert scores.dtype == torch.float32 and scores.stride(1) == 1 + assert kv_lens.shape == (scores.shape[0],) and block_tables.shape[0] == scores.shape[0] + assert select_k & (select_k - 1) == 0, f"select_k {select_k} must be a power of two" + assert out_width >= select_k * kpool + kpool - 1 + + +kpool_select = KpoolSelect() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/kpool_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/kpool_update.py new file mode 100644 index 000000000000..501208970d90 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/kpool_update.py @@ -0,0 +1,282 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pool-key refresh of the paged indexer cache, one program per row. + +Pool ``j`` of a request is its positions ``[j * kpool, (j + 1) * kpool)``. Its key sits in the +last ``head_dim`` columns of the indexer row ``[k | gate | pool key]`` of its first member and +is, per column ``d``, + + key[d] = sum_m softmax_m(gate[m, d] + ape[m, d]) * k[m, d] + +over the members ``m`` visible at the refreshing row's position (later members are masked): +an FP32 softmax, the probabilities and the products rounded to BF16, an FP32 sum and a BF16 +key, the numerics of the model's pool compression. A row refreshes the pool its position falls +in, so an incomplete pool carries the key of its visible members and the write of its last +member finalizes it. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +@triton.jit +def _kpool_update_kernel( + POOL, + BT, + POS, + APE, + REQ, + slot_stride, + row_stride, + bt_stride, + tpb, + HD: tl.constexpr, + KPOOL: tl.constexpr, + HAS_REQ: tl.constexpr, + ROWS_PER_TABLE: tl.constexpr, +): + """Recompute the pool key of the pool containing ``POS[b]``. + + Members beyond ``POS[b]`` are masked (they are not visible yet), so an + incomplete pool carries the key of its visible members and is finalized by + the write of its last member. Row ``b`` reads block table ``REQ[b]`` when + ``HAS_REQ`` (several rows of one request share a table), else table + ``b // ROWS_PER_TABLE``. With ``ROWS_PER_TABLE > 1`` the rows of one table are + consecutive positions; a row whose successor falls in the same pool is skipped, + so each pool is written once with its last visible member, as sequential + updates would leave it. Numerics follow ``build_pools``: fp32 + softmax over ``gate + ape``, bf16 probabilities and products, fp32 sum, + bf16 result. + """ + b = tl.program_id(0) + if HAS_REQ: + tbl = tl.load(REQ + b).to(tl.int64) + else: + tbl = (b // ROWS_PER_TABLE).to(tl.int64) + pos = tl.load(POS + b).to(tl.int64) + if ROWS_PER_TABLE > 1: + if b % ROWS_PER_TABLE < ROWS_PER_TABLE - 1: + if tl.load(POS + b + 1).to(tl.int64) // KPOOL == pos // KPOOL: + return + start = (pos // KPOOL) * KPOOL + m = tl.arange(0, KPOOL) + d = tl.arange(0, HD) + mem_pos = start + m + valid = mem_pos <= pos + page = mem_pos // tpb + off = mem_pos - page * tpb + slot = tl.load(BT + tbl * bt_stride + page, mask=valid, other=0).to(tl.int64) + row = POOL + slot * slot_stride + off * row_stride + k = tl.load(row[:, None] + d[None, :], mask=valid[:, None], other=0.0).to(tl.float32) + g = tl.load(row[:, None] + HD + d[None, :], mask=valid[:, None], other=0.0).to(tl.float32) + ape = tl.load(APE + m[:, None] * HD + d[None, :]).to(tl.float32) + logits = tl.where(valid[:, None], g + ape, float("-inf")) + mx = tl.max(logits, axis=0) + e = tl.exp(logits - mx[None, :]) + probs = e / tl.sum(e, axis=0)[None, :] + probs = probs.to(tl.bfloat16).to(tl.float32) + prod = (probs * k).to(tl.bfloat16).to(tl.float32) + pool_key = tl.sum(prod, axis=0) + page0 = start // tpb + off0 = start - page0 * tpb + slot0 = tl.load(BT + tbl * bt_stride + page0).to(tl.int64) + tl.store(POOL + slot0 * slot_stride + off0 * row_stride + 2 * HD + d, pool_key.to(tl.bfloat16)) + + +class KpoolUpdate(OpWrapper): + """Refresh, in ``index_pool``, the key of the pool containing each row's position. + + ``index_pool`` BF16 ``[slots, tokens_per_block, >= 3 * head_dim]``, ``block_tables`` + ``[tables, pages]``, ``ape`` BF16 ``[kpool, head_dim]`` (the compression position bias). + Two row layouts: + + - no ``request_ids``: ``positions`` ``[N, T]`` holds ``T`` consecutive positions of table + ``n`` (a verify step's draft rows), or ``[N]`` one position of table ``n`` (single-token + decode). The result equals refreshing the ``T`` columns in order: of several rows in one + pool only the last writes. + - ``request_ids`` int32 ``[N]`` with ``positions`` ``[N]``: row ``i`` reads table + ``request_ids[i]`` (packed context rows). The caller passes one row per pool, the + pool-final positions a step completes: two rows on one pool race. + + Writes only the pool-key columns of each refreshed pool's first-member row; returns None. + """ + + ARCHS = frozenset({Arch.SM_100}) + # The indexer key is 128 wide (cache rows of 3 * 128 BF16), pools of 4 over 32-token + # blocks. Generation cells give the requests' KV length after the step; request b's is + # kv - b % 4, so draft rows straddle pools differently per request. Context cells give + # (cached, new) per request: the rows are the step's pool-final positions. + CELLS: tuple[Cell, ...] = ( + Cell( + why="MTP-3 verify of one 13K-token request", + spec=dict(requests=1, tokens=4, kv=13312), + ), + Cell( + why="MTP-3 verify of eight 13K-token requests", + spec=dict(requests=8, tokens=4, kv=13312), + ), + Cell( + why="MTP-3 verify of 64 requests, the largest CUDA-graph batch", + spec=dict(requests=64, tokens=4, kv=13312), + ), + Cell( + why="single-token decode of eight 13K-token requests: 1-D positions", + spec=dict(requests=8, tokens=1, kv=13312), + ), + Cell( + why="a 13,312-token prompt: its 3,328 pool-final positions", + spec=dict(context=((0, 13312),)), + ), + Cell( + why="two requests: a chunk completing a cached pool, and a fresh prompt", + spec=dict(context=((4099, 3000), (0, 1313))), + ), + ) + + note = """ + The members of a refreshed pool are read from the `[k | gate]` columns, so a step writes + its rows' indexer state (`mla_prologue`, `mla_context_prologue`) before refreshing. + + `request_ids` with 2-D positions raises. Several rows of one table that are not + consecutive positions, or two rows on one pool with `request_ids`, leave a pool key that + depends on scheduling; `is_valid` refuses both. + """ + + def __call__( + self, + index_pool: torch.Tensor, + block_tables: torch.Tensor, + positions: torch.Tensor, + ape: torch.Tensor, + tokens_per_block: int, + *, + head_dim: int, + kpool: int, + request_ids: torch.Tensor | None = None, + ) -> None: + if request_ids is not None and positions.dim() != 1: + # The same-pool skip compares POS[b] with POS[b + 1], which is only valid + # when consecutive rows belong to one table. + raise ValueError("kpool_update: request_ids requires 1-D positions") + rows_per_table = positions.shape[1] if positions.dim() == 2 else 1 + positions = positions.reshape(-1) + n = positions.shape[0] + if n == 0: + return None + _kpool_update_kernel[(n,)]( + index_pool, + block_tables, + positions, + ape, + positions if request_ids is None else request_ids, + index_pool.stride(0), + index_pool.stride(1), + block_tables.stride(0), + tokens_per_block, + HD=head_dim, + KPOOL=kpool, + HAS_REQ=request_ids is not None, + ROWS_PER_TABLE=rows_per_table, + num_warps=4, + ) + return None + + def reference( + self, + index_pool: torch.Tensor, + block_tables: torch.Tensor, + positions: torch.Tensor, + ape: torch.Tensor, + tokens_per_block: int, + *, + head_dim: int, + kpool: int, + request_ids: torch.Tensor | None = None, + ) -> None: + """The refresh in plain torch, in place on ``index_pool``: every row's key from the + members visible at its position, and of the rows on one pool the last one's, as + refreshing the rows one after another leaves it. The BF16 roundings of the + probabilities and products are the pool key's definition, not a kernel artifact.""" + rows_per_table = positions.shape[1] if positions.dim() == 2 else 1 + pos = positions.reshape(-1).long() + n = pos.shape[0] + if n == 0: + return None + dev = pos.device + order = torch.arange(n, device=dev) + table = order // rows_per_table if request_ids is None else request_ids.long() + start = pos // kpool * kpool + member = start[:, None] + torch.arange(kpool, device=dev) # [n, kpool] + visible = member <= pos[:, None] + page = (member // tokens_per_block).clamp(max=block_tables.shape[1] - 1) + slot = block_tables[table[:, None], page].long() + off = member % tokens_per_block + rows = index_pool[slot, off].float() # [n, kpool, width] + # Members past the position are masked out of the sums, whatever their rows hold. + rows = rows.masked_fill(~visible[:, :, None], 0.0) + k, gate = rows[..., :head_dim], rows[..., head_dim : 2 * head_dim] + logits = (gate + ape.float()).masked_fill(~visible[:, :, None], float("-inf")) + probs = logits.softmax(dim=1).to(torch.bfloat16).float() + key = (probs * k).to(torch.bfloat16).float().sum(dim=1).to(torch.bfloat16) + # The last row on each pool's first-member row wins. + dest = slot[:, 0] * tokens_per_block + off[:, 0] + last = torch.full((int(dest.max()) + 1,), -1, dtype=torch.long, device=dev) + last = last.scatter_reduce(0, dest, order, reduce="amax") + keep = last[dest] == order + index_pool[slot[keep, 0], off[keep, 0], 2 * head_dim : 3 * head_dim] = key[keep] + return None + + def is_valid( + self, + index_pool: torch.Tensor, + block_tables: torch.Tensor, + positions: torch.Tensor, + ape: torch.Tensor, + tokens_per_block: int, + *, + head_dim: int, + kpool: int, + request_ids: torch.Tensor | None = None, + ) -> None: + assert index_pool.dtype == torch.bfloat16 and index_pool.dim() == 3 + assert index_pool.stride(2) == 1 and index_pool.shape[1] == tokens_per_block + assert index_pool.shape[2] >= 3 * head_dim, "rows are [k | gate | pool key]" + for name, width in (("head_dim", head_dim), ("kpool", kpool)): + assert width > 0 and width & (width - 1) == 0, f"{name} {width} must be a power of 2" + # The kernel reads ape at m * head_dim + d. + assert ape.shape == (kpool, head_dim) and ape.is_contiguous() + assert block_tables.dim() == 2 and block_tables.stride(1) == 1 + assert positions.dim() in (1, 2) + pos = positions.reshape(-1).long() + n = pos.shape[0] + if n == 0: + return + rows_per_table = positions.shape[1] if positions.dim() == 2 else 1 + if positions.dim() == 2: + # The same-pool skip lets the later of two consecutive rows write. + assert bool((positions.diff(dim=1) == 1).all()), "a table's rows must be consecutive" + if request_ids is None: + table = torch.arange(n, device=pos.device) // rows_per_table + else: + assert request_ids.shape == (n,) + table = request_ids.long() + assert 0 <= int(table.min()) and int(table.max()) < block_tables.shape[0] + assert int(pos.min()) >= 0 + assert int(pos.max()) < block_tables.shape[1] * tokens_per_block + # One writer per pool and table: with request_ids one row, else one table's rows. + first = pos // kpool * kpool + page = first // tokens_per_block + dest = block_tables[table, page].long() * tokens_per_block + first % tokens_per_block + if request_ids is not None: + assert dest.unique().numel() == n, "two rows refresh one pool" + else: + pairs = torch.stack([dest, table]).unique(dim=1) + assert pairs.shape[1] == dest.unique().numel(), "two tables share a pool's row" + + +kpool_update = KpoolUpdate() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/latent_values.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/latent_values.py new file mode 100644 index 000000000000..10abb6febb8f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/latent_values.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The absorbed value projection of a sparse MLA decode step, written row-major for o_proj. + +``out[r, h * V : (h + 1) * V] = latent[r, h, :] @ w_v_t[h]`` (BF16 inputs, FP32 +accumulation), and zero for rows that selected no key: one launch in place of a batched GEMM +that returns head-major ``[H, R, V]``, a transposing copy and a select over the result. The +target runs it on its generation rows with a rank's 16 heads (``MlaLayer.generation_partial``). +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + + +@triton.jit +def _latent_values_kernel( + lat_ptr, + w_ptr, + keep_ptr, + out_ptr, + R, + H: tl.constexpr, + L: tl.constexpr, + V: tl.constexpr, + BR: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, +): + h = tl.program_id(0) + n = tl.program_id(1) * BN + tl.arange(0, BN) + r = tl.program_id(2) * BR + tl.arange(0, BR) + rm = r < R + acc = tl.zeros((BR, BN), tl.float32) + for k0 in range(0, L, BK): + k = k0 + tl.arange(0, BK) + a = tl.load(lat_ptr + (r[:, None] * H + h) * L + k[None, :], mask=rm[:, None], other=0.0) + b = tl.load(w_ptr + (h * L + k[:, None]) * V + n[None, :]) + acc += tl.dot(a, b) + keep = tl.load(keep_ptr + r, mask=rm, other=0) != 0 + acc = tl.where(keep[:, None], acc, 0.0) + tl.store( + out_ptr + r[:, None] * (H * V) + h * V + n[None, :], acc.to(tl.bfloat16), mask=rm[:, None] + ) + + +class LatentValues(OpWrapper): + """``latent`` BF16 [R, H, L], ``w_v_t`` BF16 [H, L, V], ``keep`` bool [R] -> BF16 [R, H * V].""" + + ARCHS = frozenset({Arch.SM_100}) + CELLS: tuple[Cell, ...] = ( + Cell(why="MTP-3 verify of one request", spec=dict(r=4, h=16, l=512, v=256)), + Cell(why="MTP-3 verify of eight requests", spec=dict(r=32, h=16, l=512, v=256)), + Cell(why="CUDA-graph batch of 64 verify requests", spec=dict(r=256, h=16, l=512, v=256)), + ) + + def __call__( + self, latent: torch.Tensor, w_v_t: torch.Tensor, keep: torch.Tensor + ) -> torch.Tensor: + rows, heads, lat = latent.shape + v = w_v_t.shape[2] + out = torch.empty((rows, heads * v), dtype=torch.bfloat16, device=latent.device) + _latent_values_kernel[(heads, v // 32, triton.cdiv(rows, 16))]( + latent, w_v_t, keep, out, rows, heads, lat, v, 16, 32, 128, num_warps=4, num_stages=4 + ) + return out + + def reference( + self, latent: torch.Tensor, w_v_t: torch.Tensor, keep: torch.Tensor + ) -> torch.Tensor: + with true_fp32_matmul(): + out = torch.einsum("rhl,hlv->rhv", latent.float(), w_v_t.float()) + out = torch.where(keep[:, None, None], out, 0.0) + return out.reshape(latent.shape[0], -1).to(torch.bfloat16) + + def is_valid(self, latent: torch.Tensor, w_v_t: torch.Tensor, keep: torch.Tensor) -> None: + assert latent.dtype == torch.bfloat16 and latent.is_contiguous() and latent.dim() == 3 + rows, heads, lat = latent.shape + # The kernel reads w_v_t at (h * L + k) * V + n and takes H, L and V from latent and + # w_v_t separately: a weight of other heads or another latent width is read misaligned. + assert w_v_t.dtype == torch.bfloat16 and w_v_t.is_contiguous() and w_v_t.dim() == 3 + assert w_v_t.shape[:2] == (heads, lat), ( + f"w_v_t must be [{heads}, {lat}, V] for latent {tuple(latent.shape)}; " + f"got {tuple(w_v_t.shape)}" + ) + assert keep.dtype == torch.bool and keep.shape == (rows,) + # Whole 32-column tiles of V and 128-wide K steps, with no masks on either. + assert w_v_t.shape[2] % 32 == 0 and lat % 128 == 0 + + +latent_values = LatentValues() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/mla_context_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/mla_context_prologue.py new file mode 100644 index 000000000000..dcb4b605091a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/mla_context_prologue.py @@ -0,0 +1,281 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The per-row work between the sparse MLA input GEMMs and the query projections, for packed +context rows, in one launch. + +The context-row counterpart of ``mla_prologue``. For each packed context row ``r`` of a prefill +step: the RMSNorms of the q_a and kv_a halves of the fused input GEMM output, the LayerNorm of +the indexer key, and the paged-cache writes of the normalized latent (``clamp(latent * scale, ++-448)`` as E4M3, into the FP8 latent cache) and of the indexer row ``[k | gate]`` at cache +position ``positions[r]`` of block table ``block_tables[request_ids[r]]``. It returns the +normalized q_a rows (``q_resid``) for the q_b and indexer query projections: one launch in place +of RMSNorm, RMSNorm, LayerNorm, cat and the indexed cache writes. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +@triton.jit +def _context_prologue_kernel( + qkv_ptr, + idx_ptr, + qw_ptr, + kvw_ptr, + lnw_ptr, + lnb_ptr, + q_out_ptr, + pos_ptr, + req_ptr, + bt_ptr, + latent_pool_ptr, + index_pool_ptr, + scale_ptr, + rms_eps, + ln_eps, + qkv_stride, + idx_stride, + bt_stride, + latent_slot_stride, + latent_row_stride, + index_slot_stride, + index_row_stride, + tpb, + Q_DIM: tl.constexpr, + Q_PAD: tl.constexpr, + KV_DIM: tl.constexpr, + HEAD_DIM: tl.constexpr, +): + # A prefill step packs up to max_num_tokens rows: 64-bit row offsets. + r = tl.program_id(0).to(tl.int64) + # q_a RMSNorm + qd = tl.arange(0, Q_PAD) + qm = qd < Q_DIM + q = tl.load(qkv_ptr + r * qkv_stride + qd, mask=qm, other=0.0).to(tl.float32) + q = q * tl.rsqrt(tl.sum(q * q, axis=0) / Q_DIM + rms_eps) + q = q * tl.load(qw_ptr + qd, mask=qm, other=0.0).to(tl.float32) + tl.store(q_out_ptr + r * Q_DIM + qd, q.to(tl.bfloat16), mask=qm) + # kv_a RMSNorm -> latent cache + kd = tl.arange(0, KV_DIM) + kv = tl.load(qkv_ptr + r * qkv_stride + Q_DIM + kd).to(tl.float32) + kv = kv * tl.rsqrt(tl.sum(kv * kv, axis=0) / KV_DIM + rms_eps) + kv = (kv * tl.load(kvw_ptr + kd).to(tl.float32)).to(tl.bfloat16) + tbl = tl.load(req_ptr + r).to(tl.int64) + pos = tl.load(pos_ptr + r).to(tl.int64) + page = pos // tpb + off = pos - page * tpb + slot = tl.load(bt_ptr + tbl * bt_stride + page).to(tl.int64) + latent_dst = latent_pool_ptr + slot * latent_slot_stride + off * latent_row_stride + kd + scaled = kv.to(tl.float32) * tl.load(scale_ptr) + tl.store(latent_dst, tl.clamp(scaled, -448.0, 448.0).to(latent_pool_ptr.dtype.element_ty)) + # indexer key LayerNorm and gate -> index cache + hd = tl.arange(0, HEAD_DIM) + k = tl.load(idx_ptr + r * idx_stride + hd).to(tl.float32) + mean = tl.sum(k, axis=0) / HEAD_DIM + diff = k - mean + var = tl.sum(diff * diff, axis=0) / HEAD_DIM + k = diff * tl.rsqrt(var + ln_eps) * tl.load(lnw_ptr + hd).to(tl.float32) + tl.load( + lnb_ptr + hd + ).to(tl.float32) + gate = tl.load(idx_ptr + r * idx_stride + HEAD_DIM + hd) + index_dst = index_pool_ptr + slot * index_slot_stride + off * index_row_stride + tl.store(index_dst + hd, k.to(tl.bfloat16).to(index_pool_ptr.dtype.element_ty)) + tl.store(index_dst + HEAD_DIM + hd, gate.to(index_pool_ptr.dtype.element_ty)) + + +class MlaContextPrologue(OpWrapper): + """``q_resid`` BF16 ``[T, q_lora]`` from ``qkv`` ``[T, q_lora + kv_lora]`` and ``idx`` ``[T, + >= 2 * head_dim]`` (indexer key | gate | ...), writing the latent and indexer rows of + position ``positions[r]`` (``[T]``) of table ``block_tables[request_ids[r]]`` into the + paged pools. Only the ``[k | gate]`` columns of the indexer rows are written; the pool-key + columns are ``kpool_update``'s. + """ + + ARCHS = frozenset({Arch.SM_100}) + # GLM-5.3-Flash: q_lora 1536, kv_lora 512, an indexer key of 128 (the indexer GEMM output + # is [k | gate | 32 head weights], 288 wide), RMSNorm eps 1e-5, key LayerNorm eps 1e-6, + # 32-token blocks of an FP8 latent cache. Every rank writes every context row of a step, + # so the cells are whole steps: (cached, new) per context request, in batch order. + CELLS: tuple[Cell, ...] = ( + Cell(why="one 13,312-token prompt", spec=dict(context=((0, 13312),))), + Cell( + why="two requests: a chunk continuing an unaligned prefix, and a fresh prompt", + spec=dict(context=((4099, 3000), (0, 1313))), + ), + Cell( + why="a 37-token prompt in a mixed step: table 0 of the step's eight", + spec=dict(context=((0, 37),), tables=8), + ), + ) + + note = """ + The block table of row r is `block_tables[request_ids[r]]`, a row of the step's tables in + batch order: context requests first, so a context request's id is its index among them. + + Rows writing the same cache position race; the target's context schedule + (`ContextRows`) never packs a position twice. The latent pool is the FP8 cache this target + serves; only its write (scale, clamp, E4M3) is compiled in. + """ + + def __call__( + self, + qkv: torch.Tensor, + idx: torch.Tensor, + q_norm_weight: torch.Tensor, + kv_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + k_norm_bias: torch.Tensor, + rms_eps: float, + ln_eps: float, + positions: torch.Tensor, + request_ids: torch.Tensor, + block_tables: torch.Tensor, + latent_pool: torch.Tensor, + index_pool: torch.Tensor, + tokens_per_block: int, + scale: torch.Tensor, + ) -> torch.Tensor: + q_dim, kv_dim = q_norm_weight.shape[0], kv_norm_weight.shape[0] + head_dim = k_norm_weight.shape[0] + rows = positions.shape[0] + q_resid = torch.empty((rows, q_dim), dtype=torch.bfloat16, device=qkv.device) + if rows == 0: + return q_resid + _context_prologue_kernel[(rows,)]( + qkv, + idx, + q_norm_weight, + kv_norm_weight, + k_norm_weight, + k_norm_bias, + q_resid, + positions, + request_ids, + block_tables, + latent_pool, + index_pool, + scale, + rms_eps, + ln_eps, + qkv.stride(0), + idx.stride(0), + block_tables.stride(0), + latent_pool.stride(0), + latent_pool.stride(1), + index_pool.stride(0), + index_pool.stride(1), + tokens_per_block, + Q_DIM=q_dim, + Q_PAD=triton.next_power_of_2(q_dim), + KV_DIM=kv_dim, + HEAD_DIM=head_dim, + num_warps=4, + ) + return q_resid + + def reference( + self, + qkv: torch.Tensor, + idx: torch.Tensor, + q_norm_weight: torch.Tensor, + kv_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + k_norm_bias: torch.Tensor, + rms_eps: float, + ln_eps: float, + positions: torch.Tensor, + request_ids: torch.Tensor, + block_tables: torch.Tensor, + latent_pool: torch.Tensor, + index_pool: torch.Tensor, + tokens_per_block: int, + scale: torch.Tensor, + ) -> torch.Tensor: + """FP32 norms rounded once to BF16, then the indexed cache writes, in plain torch.""" + + def rms(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + x = x.float() + return (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + rms_eps) * w.float()).to( + torch.bfloat16 + ) + + q_dim, head_dim = q_norm_weight.shape[0], k_norm_weight.shape[0] + q_resid = rms(qkv[:, :q_dim], q_norm_weight) + if positions.shape[0] == 0: + return q_resid + latent = rms(qkv[:, q_dim:], kv_norm_weight) + k = torch.nn.functional.layer_norm( + idx[:, :head_dim].float(), + (head_dim,), + k_norm_weight.float(), + k_norm_bias.float(), + ln_eps, + ).to(torch.bfloat16) + packed = torch.cat([k, idx[:, head_dim : 2 * head_dim]], dim=-1) + pos = positions.long() + page = block_tables[request_ids.long(), pos // tokens_per_block].long() + off = pos % tokens_per_block + quantized = (latent.float() * scale).clamp(-448, 448).to(latent_pool.dtype) + # Indexed writes into FP8 payloads go through their byte views. + latent_pool.view(torch.uint8)[page, off] = quantized.view(torch.uint8) + index_pool[page, off, : 2 * head_dim] = packed.to(index_pool.dtype) + return q_resid + + def is_valid( + self, + qkv: torch.Tensor, + idx: torch.Tensor, + q_norm_weight: torch.Tensor, + kv_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + k_norm_bias: torch.Tensor, + rms_eps: float, + ln_eps: float, + positions: torch.Tensor, + request_ids: torch.Tensor, + block_tables: torch.Tensor, + latent_pool: torch.Tensor, + index_pool: torch.Tensor, + tokens_per_block: int, + scale: torch.Tensor, + ) -> None: + for t in (qkv, idx): + assert t.dtype == torch.bfloat16 and t.dim() == 2 and t.stride(-1) == 1 + kv_dim, head_dim = kv_norm_weight.shape[0], k_norm_weight.shape[0] + assert qkv.shape[1] == q_norm_weight.shape[0] + kv_dim + for n in (kv_dim, head_dim): + assert n & (n - 1) == 0, f"width {n} must be a power of two" + assert idx.shape[1] >= 2 * head_dim and k_norm_bias.shape == k_norm_weight.shape + for w in (q_norm_weight, kv_norm_weight, k_norm_weight, k_norm_bias): + assert w.is_contiguous() + # The kernel reads row r's position and table id at offset r. + assert positions.dim() == 1 and positions.is_contiguous() + assert request_ids.shape == positions.shape and request_ids.is_contiguous() + assert qkv.shape[0] == idx.shape[0] == positions.shape[0] + assert block_tables.dim() == 2 and block_tables.stride(-1) == 1 + # Only the FP8 cache's write (scale, clamp, E4M3) is compiled in. + assert latent_pool.dtype == torch.float8_e4m3fn, f"latent pool {latent_pool.dtype}" + assert latent_pool.shape[-1] == kv_dim and latent_pool.stride(-1) == 1 + assert scale.dtype == torch.float32 and scale.numel() == 1 + assert index_pool.dtype == torch.bfloat16 and index_pool.stride(-1) == 1 + assert index_pool.shape[-1] >= 2 * head_dim and tokens_per_block > 0 + if positions.shape[0] == 0: + return + req = request_ids.long() + assert 0 <= int(req.min()) and int(req.max()) < block_tables.shape[0] + pos = positions.long() + assert int(pos.min()) >= 0 + assert int(pos.max()) < block_tables.shape[1] * tokens_per_block, ( + "a position past the table width reads the next table's pages" + ) + # Two rows on one (table, position) race on the cache row. + key = req * (block_tables.shape[1] * tokens_per_block) + pos + assert key.unique().numel() == key.numel(), "a cache position is written twice" + + +mla_context_prologue = MlaContextPrologue() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/mla_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/mla_prologue.py new file mode 100644 index 000000000000..a155e45e1f05 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/mla_prologue.py @@ -0,0 +1,264 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The per-row work between the sparse MLA input GEMMs and the query projections, in one launch. + +For each generation row: the RMSNorms of the q_a and kv_a halves of the fused input GEMM +output, the LayerNorm of the indexer key, and the paged-cache writes of the normalized +latent (``clamp(latent * scale, +-448)`` as E4M3, into the FP8 latent cache) and of the +indexer row ``[k | gate]``. It returns the normalized q_a rows (``q_resid``) for the q_b and +indexer query projections: one launch in place of RMSNorm, RMSNorm, LayerNorm, cat and a +cache append. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +@triton.jit +def _prologue_kernel( + qkv_ptr, + idx_ptr, + qw_ptr, + kvw_ptr, + lnw_ptr, + lnb_ptr, + q_out_ptr, + pos_ptr, + bt_ptr, + latent_pool_ptr, + index_pool_ptr, + scale_ptr, + rms_eps, + ln_eps, + qkv_stride, + idx_stride, + bt_stride, + latent_slot_stride, + latent_row_stride, + index_slot_stride, + index_row_stride, + tpb, + ROWS_PER_TABLE: tl.constexpr, + Q_DIM: tl.constexpr, + Q_PAD: tl.constexpr, + KV_DIM: tl.constexpr, + HEAD_DIM: tl.constexpr, +): + r = tl.program_id(0) + # q_a RMSNorm + qd = tl.arange(0, Q_PAD) + qm = qd < Q_DIM + q = tl.load(qkv_ptr + r * qkv_stride + qd, mask=qm, other=0.0).to(tl.float32) + q = q * tl.rsqrt(tl.sum(q * q, axis=0) / Q_DIM + rms_eps) + q = q * tl.load(qw_ptr + qd, mask=qm, other=0.0).to(tl.float32) + tl.store(q_out_ptr + r * Q_DIM + qd, q.to(tl.bfloat16), mask=qm) + # kv_a RMSNorm -> latent cache + kd = tl.arange(0, KV_DIM) + kv = tl.load(qkv_ptr + r * qkv_stride + Q_DIM + kd).to(tl.float32) + kv = kv * tl.rsqrt(tl.sum(kv * kv, axis=0) / KV_DIM + rms_eps) + kv = (kv * tl.load(kvw_ptr + kd).to(tl.float32)).to(tl.bfloat16) + tbl = (r // ROWS_PER_TABLE).to(tl.int64) + pos = tl.load(pos_ptr + r).to(tl.int64) + page = pos // tpb + off = pos - page * tpb + slot = tl.load(bt_ptr + tbl * bt_stride + page).to(tl.int64) + latent_dst = latent_pool_ptr + slot * latent_slot_stride + off * latent_row_stride + kd + scaled = kv.to(tl.float32) * tl.load(scale_ptr) + tl.store(latent_dst, tl.clamp(scaled, -448.0, 448.0).to(latent_pool_ptr.dtype.element_ty)) + # indexer key LayerNorm and gate -> index cache + hd = tl.arange(0, HEAD_DIM) + k = tl.load(idx_ptr + r * idx_stride + hd).to(tl.float32) + mean = tl.sum(k, axis=0) / HEAD_DIM + diff = k - mean + var = tl.sum(diff * diff, axis=0) / HEAD_DIM + k = diff * tl.rsqrt(var + ln_eps) * tl.load(lnw_ptr + hd).to(tl.float32) + tl.load( + lnb_ptr + hd + ).to(tl.float32) + gate = tl.load(idx_ptr + r * idx_stride + HEAD_DIM + hd) + index_dst = index_pool_ptr + slot * index_slot_stride + off * index_row_stride + tl.store(index_dst + hd, k.to(tl.bfloat16).to(index_pool_ptr.dtype.element_ty)) + tl.store(index_dst + HEAD_DIM + hd, gate.to(index_pool_ptr.dtype.element_ty)) + + +class MlaPrologue(OpWrapper): + """``q_resid`` [R, q_lora] from ``qkv`` [R, q_lora + kv_lora] and ``idx`` [R, >= 2 * head_dim] + (indexer key | gate | ...), writing the latent and indexer rows of position ``positions`` + (``[N, T]``, table ``block_tables[n]``) into the paged pools: the FP8 latent pool + ``[slots, tokens_per_block, kv_lora]`` and the BF16 indexer pool ``[slots, tokens_per_block, + >= 2 * head_dim]``. ``scale`` is the FP32 ``[1]`` quantization scale of the latent cache.""" + + ARCHS = frozenset({Arch.SM_100}) + # The target's generation rows: T = 1 (plain decode and the MTP draft steps) or 4 (MTP-3 + # verify) rows per request, read through the int64 GLM block tables. + CELLS: tuple[Cell, ...] = ( + Cell(why="one row: an MTP draft step of one request", spec=dict(n=1, t=1)), + Cell(why="single-token decode of eight requests: one row per table", spec=dict(n=8, t=1)), + Cell(why="MTP-3 verify of one request", spec=dict(n=1, t=4)), + Cell(why="MTP-3 verify of eight requests", spec=dict(n=8, t=4)), + ) + + note = """ + `T` (`positions.shape[1]`) is a compile-time constant of the kernel: T = 1 and T = 4 compile + once each. + """ + + def __call__( + self, + qkv: torch.Tensor, + idx: torch.Tensor, + q_norm_weight: torch.Tensor, + kv_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + k_norm_bias: torch.Tensor, + rms_eps: float, + ln_eps: float, + positions: torch.Tensor, + block_tables: torch.Tensor, + latent_pool: torch.Tensor, + index_pool: torch.Tensor, + tokens_per_block: int, + scale: torch.Tensor, + ) -> torch.Tensor: + q_dim, kv_dim = q_norm_weight.shape[0], kv_norm_weight.shape[0] + head_dim = k_norm_weight.shape[0] + rows = positions.numel() + q_resid = torch.empty((rows, q_dim), dtype=torch.bfloat16, device=qkv.device) + _prologue_kernel[(rows,)]( + qkv, + idx, + q_norm_weight, + kv_norm_weight, + k_norm_weight, + k_norm_bias, + q_resid, + positions, + block_tables, + latent_pool, + index_pool, + scale, + rms_eps, + ln_eps, + qkv.stride(0), + idx.stride(0), + block_tables.stride(0), + latent_pool.stride(0), + latent_pool.stride(1), + index_pool.stride(0), + index_pool.stride(1), + tokens_per_block, + ROWS_PER_TABLE=positions.shape[1], + Q_DIM=q_dim, + Q_PAD=triton.next_power_of_2(q_dim), + KV_DIM=kv_dim, + HEAD_DIM=head_dim, + num_warps=4, + ) + return q_resid + + def reference( + self, + qkv: torch.Tensor, + idx: torch.Tensor, + q_norm_weight: torch.Tensor, + kv_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + k_norm_bias: torch.Tensor, + rms_eps: float, + ln_eps: float, + positions: torch.Tensor, + block_tables: torch.Tensor, + latent_pool: torch.Tensor, + index_pool: torch.Tensor, + tokens_per_block: int, + scale: torch.Tensor, + ) -> torch.Tensor: + def rms(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + x = x.float() + return (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + rms_eps) * w.float()).to( + torch.bfloat16 + ) + + q_dim, head_dim = q_norm_weight.shape[0], k_norm_weight.shape[0] + q_resid = rms(qkv[:, :q_dim], q_norm_weight) + latent = rms(qkv[:, q_dim:], kv_norm_weight) + k = torch.nn.functional.layer_norm( + idx[:, :head_dim].float(), + (head_dim,), + k_norm_weight.float(), + k_norm_bias.float(), + ln_eps, + ).to(torch.bfloat16) + packed = torch.cat([k, idx[:, head_dim : 2 * head_dim]], dim=-1) + tokens = positions.shape[1] + for r in range(positions.numel()): + pos = int(positions.view(-1)[r]) + slot = int(block_tables[r // tokens, pos // tokens_per_block]) + off = pos % tokens_per_block + latent_pool[slot, off] = ( + (latent[r].float() * scale).clamp(-448, 448).to(latent_pool.dtype) + ) + index_pool[slot, off, : 2 * head_dim] = packed[r].to(index_pool.dtype) + return q_resid + + def is_valid( + self, + qkv: torch.Tensor, + idx: torch.Tensor, + q_norm_weight: torch.Tensor, + kv_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + k_norm_bias: torch.Tensor, + rms_eps: float, + ln_eps: float, + positions: torch.Tensor, + block_tables: torch.Tensor, + latent_pool: torch.Tensor, + index_pool: torch.Tensor, + tokens_per_block: int, + scale: torch.Tensor, + ) -> None: + for t in (qkv, idx): + assert t.dtype == torch.bfloat16 and t.stride(-1) == 1 + kv_dim, head_dim = kv_norm_weight.shape[0], k_norm_weight.shape[0] + assert qkv.shape[1] == q_norm_weight.shape[0] + kv_dim + for n in (kv_dim, head_dim): + assert n & (n - 1) == 0, f"width {n} must be a power of two" + assert idx.shape[1] >= 2 * head_dim and k_norm_bias.shape == k_norm_weight.shape + # The kernel reads position r at offset r and the table of row r at row r // T. + assert positions.dim() == 2 and positions.is_contiguous() + assert not positions.is_floating_point(), f"positions are integers; got {positions.dtype}" + assert qkv.shape[0] == idx.shape[0] == positions.numel() + assert block_tables.dim() == 2 and block_tables.stride(-1) == 1 + assert block_tables.shape[0] >= positions.shape[0] + # Only the FP8 cache's write (scale, clamp, E4M3) is compiled in. + assert latent_pool.dtype == torch.float8_e4m3fn, f"latent pool {latent_pool.dtype}" + assert scale.dtype == torch.float32 and scale.numel() == 1 + assert latent_pool.dim() == 3 and latent_pool.shape[1:] == (tokens_per_block, kv_dim) + assert latent_pool.stride(-1) == 1 + assert index_pool.dtype == torch.bfloat16 and index_pool.dim() == 3 + assert index_pool.shape[1] == tokens_per_block and index_pool.stride(-1) == 1 + assert index_pool.shape[-1] >= 2 * head_dim and tokens_per_block > 0 + if positions.numel() == 0: + return + pos = positions.long() + # A position past the table's width reads the next table's pages as its own. + assert int(pos.min()) >= 0 + assert int(pos.max()) < block_tables.shape[1] * tokens_per_block, ( + f"positions up to {int(pos.max())} exceed the tables' " + f"{block_tables.shape[1] * tokens_per_block} positions" + ) + # Two rows on one cache position race on the cache row. + assert bool((pos.sort(dim=1).values.diff(dim=1) != 0).all()), "a position is written twice" + slots = block_tables[: positions.shape[0]].gather(1, pos // tokens_per_block).long() + # A slot outside the pools is written, not refused. + assert int(slots.min()) >= 0 and int(slots.max()) < min( + latent_pool.shape[0], index_pool.shape[0] + ), "a table entry addresses a slot past the pools" + + +mla_prologue = MlaPrologue() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/sparse_mla_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/sparse_mla_decode.py new file mode 100644 index 000000000000..92a05c7b5aad --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/attention/sparse_mla_decode.py @@ -0,0 +1,226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Sparse MLA over the FP8 latent cache: FlashInfer's TRTLLM-GEN native NoPE MLA decode kernel. + +Every query row ``r`` (a decode, verify or context row alike) is its own sequence of one token: +it attends, with all its heads, to the latent rows ``packed[r, :lengths[r]]`` of the cache, +which serve as both keys and values. The scores are ``q . k * bmm1_scale`` over the 512-wide +latent space (the queries are absorbed; there is no rotary part), the softmax runs in FP32, and +the output ``softmax @ v * bmm2_scale`` is BF16. For an FP8 cache with dequantization scale +``s`` and the model's softmax scale ``qk_nope_head_dim ** -0.5``, ``bmm1_scale = s * +qk_nope_head_dim ** -0.5`` and ``bmm2_scale = s``. + +Buffers the caller owns and passes on every call: + +- ``workspace``: uint8, at least ``WORKSPACE_BYTES``, zeroed before its first use. The kernel + keeps the partial results of a row split across CTAs there. +- ``counter``: uint8, contiguous, 16-byte aligned, zeroed once, at least + ``counter_bytes(rows, heads, sm_count)`` for the call's rows and heads. The kernel resets + its semaphores after every launch, so ordered launches on one stream (and replays of one + CUDA graph) share it; launches that can run concurrently each need their own. + +A CUDA-graph capture needs both buffers allocated before it starts; the target's +``SparseAttentionScratch`` (``layers/mla.py``) holds one of each for every sparse MLA layer. +""" + +from __future__ import annotations + +import torch +from flashinfer.mla import trtllm_batch_decode_with_kv_cache_mla +from flashinfer.utils import get_trtllm_gen_multi_ctas_kv_counter_bytes + +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 + +#: Workspace bytes the kernel needs: 32 MiB, the TRTLLM-GEN share of TensorRT-LLM's generation +#: attention workspace (``CUBLAS_WORKSPACE_SIZE``). +WORKSPACE_BYTES = 32 * 1024 * 1024 + +#: The native NoPE MLA geometry: absorbed 512-wide queries and latents, no rotary part, and +#: the latent cache read as 32-row pages. +KV_LORA_RANK = 512 +QK_NOPE_HEAD_DIM = 256 +PAGE_ROWS = 32 + + +def counter_bytes(rows: int, num_heads: int, sm_count: int) -> int: + """Bytes of multi-CTA semaphores a call of ``rows`` query rows of ``num_heads`` heads needs: + one 4-byte counter per (row, head), at least one per SM, rounded up to 8 counters.""" + return get_trtllm_gen_multi_ctas_kv_counter_bytes(rows, num_heads, sm_count) + + +class SparseMlaDecode(OpWrapper): + """``query`` E4M3 ``[R, H, 512]`` over ``kv_rows`` E4M3 ``[N, 1, 512]`` -> BF16 ``[R, H, 512]``. + + ``packed`` int32 ``[R, W]`` holds row ``r``'s latent row ids (rows of ``kv_rows``) in its + first ``lengths[r]`` columns and ``-1`` after them (``compact_sparse_rows``); ``lengths`` + int32 ``[R]`` is at least 1. A row whose selection was empty carries one dummy id and + length 1, and gets that row's value: the caller discards it. ``bmm1_scale`` and + ``bmm2_scale`` are Python floats (fixed in a captured graph) or FP32 device tensors of one + element. + """ + + ARCHS = frozenset({Arch.SM_100}) + # TP4 decode and verify rows attend with the rank's 16 heads; the context rows of a prefill + # step attend sequence parallel, all 64 heads for the rank's ceil(T / 4) rows, whatever the + # prompt length. Every row selects up to 512 pools of 4 plus a tail of up to 3: lengths up + # to 2,051 in packed rows of 2,112. `kv` is the requests' KV length after the step + # (generation cells); context cells give the prompt and the rank's row range. + CELLS: tuple[Cell, ...] = ( + Cell( + why="MTP-3 verify of one 13K-token request, 16 heads", + spec=dict(heads=16, requests=1, tokens=4, kv=13312), + ), + Cell( + why="MTP-3 verify of eight 13K-token requests, 16 heads", + spec=dict(heads=16, requests=8, tokens=4, kv=13312), + ), + Cell( + why="MTP-3 verify of 64 requests, the largest CUDA-graph batch; one row empty", + spec=dict(heads=16, requests=64, tokens=4, kv=13312, empty_rows=1), + ), + Cell( + why="single-token decode of eight 13K-token requests, 16 heads", + spec=dict(heads=16, requests=8, tokens=1, kv=13312), + ), + Cell( + why="one row: an MTP draft step of one 13K-token request, 16 heads", + spec=dict(heads=16, requests=1, tokens=1, kv=13312), + ), + Cell( + why="one rank's sequence-parallel quarter of a 13,312-token prompt, 64 heads", + spec=dict(heads=64, context=((0, 13312),), rows=(0, 3328)), + ), + Cell( + why="rank 1's quarter (rows 10-19) of a 37-token prompt, 64 heads: short rows", + spec=dict(heads=64, context=((0, 37),), rows=(10, 20)), + ), + ) + + note = """ + The kernel rounds the softmax probabilities to E4M3 for the value product (the queries + and the cache are E4M3 already), so an output moves by a few percent of its own scale + against FP32 attention; the entry's test derives its band from that rounding. + + The query is E4M3: the target converts the absorbed BF16 query with `saturate_fp8` + (`clamp(+-448)`, then E4M3). `kv_rows` is the uniformly strided row view of the latent pool + (row id = base_row + slot * rows_per_slot + offset); the kernel reads it as 32-row pages, + so its row count is a multiple of 32. + + Heads: the target attends with 16 (generation rows, one rank's share of 64) and 64 + (context rows); TensorRT-LLM's sparse MLA attention also runs this kernel at 32, and + 64 < H < 128 has no TRTLLM-GEN kernel and raises. A zero length is not supported by the + kernel; `is_valid` refuses it. The launch uses programmatic dependent launch unless + `TRTLLM_ENABLE_PDL=0`. + """ + + def __call__( + self, + query: torch.Tensor, + kv_rows: torch.Tensor, + packed: torch.Tensor, + lengths: torch.Tensor, + workspace: torch.Tensor, + counter: torch.Tensor, + bmm1_scale: float | torch.Tensor, + bmm2_scale: float | torch.Tensor, + ) -> torch.Tensor: + return trtllm_batch_decode_with_kv_cache_mla( + query=query.unsqueeze(1), + kv_cache=kv_rows.view(-1, 1, PAGE_ROWS, KV_LORA_RANK), + workspace_buffer=workspace, + qk_nope_head_dim=QK_NOPE_HEAD_DIM, + kv_lora_rank=KV_LORA_RANK, + qk_rope_head_dim=0, + block_tables=packed.unsqueeze(1), + seq_lens=lengths, + max_seq_len=packed.shape[1], + sparse_mla_top_k=packed.shape[1], + bmm1_scale=bmm1_scale, + bmm2_scale=bmm2_scale, + enable_pdl=get_env_enable_pdl(), + backend="trtllm-gen", + sparse_mla_top_k_lens=lengths, + multi_ctas_kv_counter_buffer=counter, + use_fp16_softmax=False, + ).squeeze(1) + + def reference( + self, + query: torch.Tensor, + kv_rows: torch.Tensor, + packed: torch.Tensor, + lengths: torch.Tensor, + workspace: torch.Tensor, + counter: torch.Tensor, + bmm1_scale: float | torch.Tensor, + bmm2_scale: float | torch.Tensor, + ) -> torch.Tensor: + """FP32 attention of the dequantized (exact) E4M3 operands over each row's selected + latent rows, one BF16 rounding. ``workspace`` and ``counter`` are not read.""" + rows, heads, dim = query.shape + # Gathered as bytes: advanced indexing does not take every FP8 dtype. + latent = kv_rows.reshape(-1, dim).view(torch.uint8) + out = torch.empty(rows, heads, dim, dtype=torch.bfloat16, device=query.device) + col = torch.arange(packed.shape[1], device=query.device) + for r0 in range(0, rows, 16): + part = slice(r0, min(r0 + 16, rows)) + selected = col[None, :] < lengths[part, None].long() + ids = packed[part].long().clamp(min=0) + values = latent[ids].view(kv_rows.dtype).float() # [c, W, dim] + # The padding past each length is never read: no value of it may reach the sums. + values = values.masked_fill(~selected[:, :, None], 0.0) + with true_fp32_matmul(): + logits = torch.einsum("rhd,rwd->rhw", query[part].float(), values) * bmm1_scale + probs = logits.masked_fill(~selected[:, None, :], float("-inf")).softmax(dim=-1) + with true_fp32_matmul(): + attended = torch.einsum("rhw,rwd->rhd", probs, values) * bmm2_scale + out[part] = attended.to(torch.bfloat16) + return out + + def is_valid( + self, + query: torch.Tensor, + kv_rows: torch.Tensor, + packed: torch.Tensor, + lengths: torch.Tensor, + workspace: torch.Tensor, + counter: torch.Tensor, + bmm1_scale: float | torch.Tensor, + bmm2_scale: float | torch.Tensor, + ) -> None: + assert query.dtype == kv_rows.dtype == torch.float8_e4m3fn, ( + f"the kernel is the E4M3 x E4M3 one; got query {query.dtype}, cache {kv_rows.dtype}" + ) + assert query.dim() == 3 and query.shape[2] == KV_LORA_RANK and query.is_contiguous() + rows, heads, _ = query.shape + assert heads in (16, 32, 64), f"{heads} heads; the kernel is run at 16, 32 or 64" + assert kv_rows.shape[-1] == KV_LORA_RANK and kv_rows.shape[0] % PAGE_ROWS == 0 + assert packed.dtype == torch.int32 and packed.dim() == 2 and packed.is_contiguous() + assert lengths.dtype == torch.int32 and lengths.shape == (rows,) + assert workspace.dtype == torch.uint8 and workspace.is_contiguous() + # The kernel's split-KV scratch is "the rest of the workspace": a short one is + # overrun, not refused. + assert workspace.numel() >= WORKSPACE_BYTES, ( + f"workspace of {workspace.numel()} bytes < {WORKSPACE_BYTES}" + ) + # Semaphores left nonzero by anything but this kernel desynchronize the split rows. + assert not bool(counter.any()), "the multi-CTA counter must be zero between launches" + if rows == 0: + return + # A zero length is unsupported by the kernel; a length past the row reads -1 padding. + assert int(lengths.min()) >= 1 and int(lengths.max()) <= packed.shape[1] + col = torch.arange(packed.shape[1], device=packed.device) + selected = packed[col[None, :] < lengths[:, None].long()] + # An id outside the cache is read, not refused. + assert int(selected.min()) >= 0 and int(selected.max()) < kv_rows.shape[0], ( + "every selected id must be a row of kv_rows" + ) + + +sparse_mla_decode = SparseMlaDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/comm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/comm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/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/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/comm/mhc_ar_boundary.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/comm/mhc_ar_boundary.py new file mode 100644 index 000000000000..7084775a6b3c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/comm/mhc_ar_boundary.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The tensor-parallel all-reduce of a decode step's partial rows fused with the mHC boundary after it. + +One cluster of 16 CTAs runs TensorRT-LLM's Lamport one-shot all-reduce on the TP group's all-reduce +workspace, so its calls interleave with the library's one-shot all-reduce on that workspace, and then +the boundary between two mHC-wrapped blocks: the previous block's post-mapping into the residual +streams and the next block's pre-mapping (mixer, Sinkhorn-normalized stream combination, RMS-normalized +layer input). It takes decode steps of at most 8 rows in place of the library all-reduce and +``mHC.fused_hc``. The CTA's slice of the mixer weights is loaded before the grid-dependency wait, and the 16 CTAs +add their partial mixer sums over distributed shared memory in a fixed order. The kernel is in +``csrc/mhcArBoundary.cu``. +""" + +from __future__ import annotations + +from typing import Sequence + +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 ... import native_ops +from ..mhc.mhc_big_fuse import sinkhorn + +WORLD = 4 +MAX_ROWS = 8 +MULT = 4 +HIDDEN = 4096 + + +class MhcArBoundary(OpWrapper): + """The mHC boundary after a block whose output each of the 4 ranks holds a partial of. + + ``partial`` is this rank's ``[M, 4096]`` BF16 share of the block output; their sum over the group, + rounded to BF16 like the library's one-shot all-reduce, is the block output ``x``. ``residual`` + ``[M, 4, 4096]`` BF16 holds the streams, ``post`` ``[M, 4]`` and ``comb`` ``[M, 4, 4]`` FP32 the + previous block's coefficients; ``fn``, ``scale`` and ``base`` are the next block's mixer and + ``norm_weight`` its input RMSNorm. ``workspace`` is the group's all-reduce workspace and ``rank`` + this rank in it. + + Returns ``(residual_out, post_out, comb_out, layer_input)``: the streams ``post * x + comb^T + streams`` (BF16), the next block's ``post`` and Sinkhorn-normalized ``comb`` (FP32, ``[M, 4]`` and + ``[M, 4, 4]``) and its RMS-normalized layer input (BF16) -- what the library all-reduce followed by + ``mHC.fused_hc(..., norm_weight=...)`` returns, ``post`` without its trailing unit dimension. + """ + + ARCHS = frozenset({Arch.SM_100}) + + # GLM-5.3-Flash at tp4: hc_mult 4, hidden 4096, 20 Sinkhorn iterations. + _GLM = dict(world=WORLD, mult=MULT, hidden=HIDDEN, sinkhorn_iters=20) + CELLS: tuple[Cell, ...] = ( + Cell(why="decode of one request without MTP", spec=dict(rows=1, **_GLM)), + Cell(why="MTP-3 verify of one request", spec=dict(rows=4, **_GLM)), + Cell(why="the kernel's 8-row bound", spec=dict(rows=8, **_GLM)), + ) + + note = """ + The kernel is compiled for 4 ranks, hc_mult 4 and hidden 4096 and holds at most 8 rows. The model + sends it the steps the decode MoE kernel takes (at most 4 rows); larger steps take the library + all-reduce and ``mHC.fused_hc``. It reads and rotates the workspace's + Lamport state (receive buffers, flag, counter, clear size), so its calls and the library one-shot + all-reduce's on that workspace must be ordered on one stream. The mixer products run on TF32 + tensor cores (the BF16 streams are exact in TF32) and the layer input's RMS comes from the 4 x 4 + Gram matrix of the new streams, so no second reduction over the hidden dimension is needed. + """ + + def __call__( + self, + partial: torch.Tensor, + workspace: torch.Tensor, + rank: int, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + norm_weight: torch.Tensor, + sinkhorn_iters: int, + norm_eps: float, + hc_eps: float, + sinkhorn_eps: float, + post_mult: float, + layer_input_eps: float, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + rows = partial.shape[0] + residual_out = torch.empty_like(residual) + post_out = torch.empty((rows, MULT), dtype=torch.float32, device=partial.device) + comb_out = torch.empty((rows, MULT, MULT), dtype=torch.float32, device=partial.device) + layer_input = torch.empty_like(partial) + native_ops.mhc_ar_boundary( + partial, + workspace, + rank, + residual, + post, + comb, + fn, + scale, + base, + norm_weight, + sinkhorn_iters, + norm_eps, + hc_eps, + sinkhorn_eps, + post_mult, + layer_input_eps, + residual_out, + post_out, + comb_out, + layer_input, + get_env_enable_pdl(), + ) + return residual_out, post_out, comb_out, layer_input + + def reference( + self, + partial: Sequence[torch.Tensor], + workspace: torch.Tensor, + rank: int, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + norm_weight: torch.Tensor, + sinkhorn_iters: int, + norm_eps: float, + hc_eps: float, + sinkhorn_eps: float, + post_mult: float, + layer_input_eps: float, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """The same outputs in FP32 torch. A collective: ``partial`` is the group's partials in rank + order, not this rank's one; ``workspace`` and ``rank`` are ignored. + + Rounded to BF16 where the op stores a tensor: the all-reduce result and the new streams (the + next block reads the stored streams); everything else stays FP32, with FP32 matmuls. The + collapsed layer input is normalized in FP32: the kernel takes its RMS from the Gram matrix of + the new streams and never stores the row, so unlike the unfused path (the library all-reduce, + then ``mHC.fused_hc``) it is not rounded to BF16 before the RMSNorm. That moves the layer + input by under 2 ulp per element (0.8 ulp RMS) against the unfused path. + """ + rows = residual.shape[0] + with true_fp32_matmul(): + x = torch.stack([p.float() for p in partial]).sum(0).to(torch.bfloat16).float() + streams = residual.float().reshape(rows, MULT, HIDDEN) + post_prev = post.float().reshape(rows, MULT, 1) + comb_prev = comb.float().reshape(rows, MULT, MULT) + new = (x.unsqueeze(1) * post_prev + torch.bmm(comb_prev.mT, streams)).to(torch.bfloat16) + flat = new.float().reshape(rows, MULT * HIDDEN) + rstd = torch.rsqrt(flat.square().mean(-1, keepdim=True) + norm_eps) + mixes = flat @ fn.float().T * rstd + scales = torch.cat( + [scale[0].expand(MULT), scale[1].expand(MULT), scale[2].expand(MULT * MULT)] + ).float() + mixes = mixes * scales + base.float() + pre = torch.sigmoid(mixes[:, :MULT]) + hc_eps + post_out = torch.sigmoid(mixes[:, MULT : 2 * MULT]) * post_mult + comb_out = sinkhorn( + mixes[:, 2 * MULT :].reshape(rows, MULT, MULT), sinkhorn_iters, sinkhorn_eps + ) + collapsed = (new.float() * pre.unsqueeze(-1)).sum(1) + layer_input = collapsed * torch.rsqrt( + collapsed.square().mean(-1, keepdim=True) + layer_input_eps + ) + return new, post_out, comb_out, (layer_input * norm_weight.float()).to(torch.bfloat16) + + def is_valid( + self, + partial: torch.Tensor, + workspace: torch.Tensor, + rank: int, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + norm_weight: torch.Tensor, + sinkhorn_iters: int, + norm_eps: float, + hc_eps: float, + sinkhorn_eps: float, + post_mult: float, + layer_input_eps: float, + ) -> None: + rows, hidden = partial.shape + assert 1 <= rows <= MAX_ROWS, f"{rows} rows; at most {MAX_ROWS}" + assert hidden == HIDDEN and residual.shape == (rows, MULT, HIDDEN), ( + "the kernel is compiled for hc_mult 4 and hidden 4096" + ) + assert workspace.numel() == 3 * WORLD + 2 and 0 <= rank < WORLD, ( + f"the kernel is compiled for {WORLD} ranks (a workspace of {3 * WORLD + 2} pointers)" + ) + for t in (partial, residual, norm_weight): + assert t.dtype == torch.bfloat16 and t.is_contiguous() + for t in (post, comb, fn, scale, base): + assert t.dtype == torch.float32 and t.is_contiguous() + assert post.numel() == rows * MULT and comb.numel() == rows * MULT * MULT + assert fn.shape == ((2 + MULT) * MULT, MULT * HIDDEN) + + +mhc_ar_boundary = MhcArBoundary() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/gemm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/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/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/gemm/dsv3_router_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/gemm/dsv3_router_gemm.py new file mode 100644 index 000000000000..94799a9f94d4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/gemm/dsv3_router_gemm.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MoE router logits: BF16 hidden rows times the column-major BF16 router weight, FP32 out, via +the trtllm dsv3_router_gemm_op (a small-M kernel for the router shapes it is built for, cuBLASLt +otherwise).""" + +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 + +#: GLM-5.3-Flash's router: hidden 4096 -> 288 routed experts, replicated on every rank. +GLM_ROUTER = dict(hidden=4096, experts=288) + + +class Dsv3RouterGemm(OpWrapper): + """`mat_a @ mat_b (+ bias)`, accumulated in fp32, one `dsv3_router_gemm_op` call. + + `mat_a` is `[M, K]` row-major, `mat_b` is `[K, N]` read column-major -- the stored + `[N, K]` router weight's `.t()`, so no transpose kernel runs. GLM-5.3-Flash calls it with + BF16 `[M, 4096]` hidden rows, its BF16 `[288, 4096]` router weight `.t()`, no bias and + `out_dtype=torch.float32`: FP32 `[M, 288]` logits, the input of the in-kernel routing of + `fp8_block_scale_moe_runner`. + + Two kernels behind one call, both fp32-accumulated over exact bf16 products: + + * `1 <= M <= 16`, BF16 operands, FP32 out, no bias, and `(K, N)` one of the shapes the op + is built for -- `(4096, 288)` is GLM-5.3-Flash's -- run a small-M kernel + (`invokeRouterGemm`): one 128-thread CTA per expert column, each thread summing 32 + products in order with IEEE fp32 adds, then a warp butterfly and a 4-warp sum. + * Everything else falls back to cuBLASLt (`cublas_mm_out`, compute type 32F) on the + tensor cores, whose fp32 accumulation is coarser than IEEE round-to-nearest -- tens of + fp32 ulp of the row scale at GLM's shape rather than about one: at GLM's shape, every + step of more than 16 rows, and any call with a bias or a non-FP32 output. + """ + + ARCHS = frozenset({Arch.SM_100}) + + # Only GLM-5.3-Flash calls this entry. Its own cells are the boundary between the two + # kernels at GLM's router shape; the target's row counts are in the test. + CELLS: tuple[Cell, ...] = ( + Cell( + why="16 rows: the small-M kernel's last row count", + spec=dict(rows=16, **GLM_ROUTER), + ), + Cell( + why="17 rows: the first row count on the cuBLASLt fallback", + spec=dict(rows=17, **GLM_ROUTER), + ), + ) + + note = """ + Source: compiled-in C++, `torch.ops.trtllm.dsv3_router_gemm_op` (thop/dsv3RouterGemmOp.cpp). + The small-M kernel is instantiated per row count (1..16) and per built-in `(K, N)`; the + `(4096, 288)` instantiation is GLM-5.3-Flash's. Above 16 rows, or for any other shape, + dtype, output dtype or a bias, the op runs cuBLASLt instead -- slower at small M, and + accumulated on the tensor cores. The two kernels' logits are not bitwise equal; both stay + within a few times 1e-5 of the row scale (the test's band), far below what a bf16 + rounding of the logits would move. + + Layout handling is the op's own and loud: `mat_a` must have a unit column stride and + `mat_b` a unit row stride (it raises otherwise); an operand with padded rows or a + misaligned start is copied into a dense 16-byte-aligned layout before either kernel + reads it. So no layout guard is needed. + + The output is a new `[M, N]` tensor in `out_dtype` (default `mat_a.dtype`). For the + routing that consumes it the FP32 output is what matters: rounding the logits to bf16 can + reorder near-equal experts. + + Not certified: a bias (it goes to the same cuBLASLt call as `cublas_mm`, whose hazards + `is_valid` repeats), outputs other than FP32, operand dtypes other than BF16, router + shapes other than GLM-5.3-Flash's, CUDA-graph capture, sm_103. + """ + + def __call__( + self, + mat_a: torch.Tensor, + mat_b: torch.Tensor, + bias: Optional[torch.Tensor] = None, + out_dtype: Optional[torch.dtype] = None, + ) -> torch.Tensor: + return torch.ops.trtllm.dsv3_router_gemm_op(mat_a, mat_b, bias, out_dtype) + + def reference( + self, + mat_a: torch.Tensor, + mat_b: torch.Tensor, + bias: Optional[torch.Tensor] = None, + out_dtype: Optional[torch.dtype] = None, + ) -> torch.Tensor: + """Accumulated in float64 and rounded once to `out_dtype`. + + The op's own output is FP32, so an fp32-accumulated reference would be no more + accurate than the op and the test's band would measure the reference. In float64 the + bf16 products and their 4096-term sums are exact to well below one fp32 ulp, and no + TF32 setting applies. + """ + ref = mat_a.double() @ mat_b.double() + if bias is not None: + ref = ref + bias.double() + return ref.to(out_dtype if out_dtype is not None else mat_a.dtype) + + def is_valid( + self, + mat_a: torch.Tensor, + mat_b: torch.Tensor, + bias: Optional[torch.Tensor] = None, + out_dtype: Optional[torch.dtype] = None, + ) -> None: + """Only the bias can be wrong without a loud error. + + Every layout mistake the op either rejects or repairs (see the note). A bias sends the + call to the cuBLASLt epilogue that `cublas_mm` runs, which takes a bias of the wrong + dtype or shape and returns a wrong answer, and ignores a bias on fp32 operands; those + conditions are repeated here from that entry. GLM-5.3-Flash passes no bias. + """ + if bias is None: + return + out_dt = out_dtype if out_dtype is not None else mat_a.dtype + assert mat_a.dtype != torch.float32, "bias is silently ignored for fp32 inputs" + assert bias.dtype == out_dt, f"bias dtype {bias.dtype} must equal the output dtype {out_dt}" + assert bias.shape == (mat_b.shape[1],) and bias.is_contiguous(), ( + "bias must be a contiguous [N] tensor" + ) + + +dsv3_router_gemm = Dsv3RouterGemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/gemm/fp8_blockwise_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/gemm/fp8_blockwise_gemm.py new file mode 100644 index 000000000000..82e4fcfd4b13 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/gemm/fp8_blockwise_gemm.py @@ -0,0 +1,198 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Block-FP8 linear: e4m3 activations with 1x128 FP32 scales times an e4m3 weight with 128x128 +FP32 scales, BF16 out, via the CuTe DSL Blackwell blockwise GEMM +(`cute_dsl_fp8_gemm_blackwell`).""" + +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, +) + +#: Block edge on every scale axis: 1x128 activations, 128x128 weights. +BLOCK = 128 + + +def _ceil_blocks(n: int) -> int: + return (n + BLOCK - 1) // BLOCK + + +def dequantize_1x128(data: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + """`[M, K]` fp32 activations from e4m3 codes and K-block-major `[ceil(K/128), M]` scales.""" + k = data.shape[1] + return data.float() * scales.float().t().repeat_interleave(BLOCK, dim=1)[:, :k] + + +def dequantize_128x128(codes: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + """`[N, K]` fp32 weight from e4m3 codes and `[ceil(N/128), ceil(K/128)]` scales.""" + n, k = codes.shape + expanded = scales.float().repeat_interleave(BLOCK, dim=0)[:n] + return codes.float() * expanded.repeat_interleave(BLOCK, dim=1)[:, :k] + + +#: GLM-5.3-Flash's block-FP8 FFN projections at TP4, as `(N, K)` of the stored weight: +#: hidden 4096, dense intermediate 12288 / 4, shared-expert intermediate 2048 / 4. +GLM_DENSE_GATE_UP = dict(n=2 * 3072, k=4096) +GLM_DENSE_DOWN = dict(n=4096, k=3072) +GLM_SHARED_GATE_UP = dict(n=2 * 512, k=4096) +GLM_SHARED_DOWN = dict(n=4096, k=512) + + +class Fp8BlockwiseGemm(OpWrapper): + """`(input * input_scale) @ (weight * weight_scale).T`, fp32-accumulated, bf16 out. + + `input` is `[M, K]` e4m3 and `input_scale` its `[ceil(K/128), M]` fp32 scales -- + K-block-major, exactly what `fp8_quantize_1x128` returns on the sm_100 family. + `weight` is the stored `[N, K]` e4m3 weight and `weight_scale` its `[ceil(N/128), + ceil(K/128)]` fp32 block scales -- a block-FP8 checkpoint's `weight_scale_inv` as stored, + not resmoothed to UE8M0 powers of two. Returns a new `[M, N]` bf16 tensor. + + The kernel multiplies each 128-wide K block's fp32 MMA partial by its activation and + weight scales and accumulates the scaled partials in fp32. That is the block-FP8 `Linear` + of a checkpoint quantized with 128x128 weight blocks on SM100/SM103 when the DeepGEMM + path is off (`fp8_quantize_1x128` then this op). + + Fusion boundary: the GEMM and its dequantization only. No bias, no activation, no + quantization of the input -- the caller quantizes the bf16 rows with + `fp8_quantize_1x128` and may reuse one quantization for several GEMMs that read the same + rows. + """ + + ARCHS = frozenset({Arch.SM_100}) + + # Only GLM-5.3-Flash calls this entry. Its own cells are the M-tile edges at GLM's widest + # FFN projection; the projections and row counts the target runs are in the test. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one row: the decode row count, a single partial M tile", + spec=dict(m=1, **GLM_DENSE_GATE_UP), + ), + Cell( + why="200 rows: one full and one partial 128-row M tile", + spec=dict(m=200, **GLM_DENSE_GATE_UP), + ), + ) + + note = """ + Source: the CuTe DSL kernel `Sm100BlockwiseGemmKernel` + (`_torch/kernels/blackwell/blockwise_gemm/blockwise_gemm.py`) behind the + `trtllm::cute_dsl_fp8_gemm_blackwell` custom op (`_torch/custom_ops/cute_dsl_custom_ops.py`), + JIT-compiled per tactic on first use and autotuned by `AutoTuner`; outside an `autotune()` + context a cold cache runs the default tactic (1-CTA, 128x128 MMA tile, 1x1 cluster). The + op raises off the sm_100 family and for any `output_dtype` but bf16. `use_tvm_ffi` + changes the launch path, not the result. + + The kernel reads all four operands through raw data pointers with fixed dense layouts -- + `input` `[M, K]` and `weight` `[N, K]` K-contiguous, `input_scale` `[ceil(K/128), M]` with + M contiguous, `weight_scale` `[ceil(N/128), ceil(K/128)]` row-major -- each assumed + 16-byte aligned, and takes K from `weight`. Strides, dtypes and the operands' shapes are + not checked against each other; `is_valid` states what has to hold. + + At `M == K/128` (32 rows at K 4096, 24 at 3072, 12 at 1536, 4 at 512) a token-major + `[M, K/128]` activation scale has the shape of the K-block-major one, so no check can tell + them apart: pass the scales `fp8_quantize_1x128` returns, as returned. A row slice + `x_fp8[a:b]` needs `x_scale[:, a:b].contiguous()`, not the strided view. + + Certified: GLM-5.3-Flash's block-FP8 projections at TP4 -- the dense MLP (6144 x 4096 and + 4096 x 3072), the shared expert (1024 x 4096 and 4096 x 512) and the sequence-parallel + MLA prefill's q_b and o (16384 x 1536, 4096 x 16384) -- from 1 to 13,312 rows, all with N + and K multiples of 128. + + The target's decode steps run this op, after `fp8_quantize_1x128`, inside CUDA graphs at + the tactics its warmup tuned; the test certifies that at one decode shape (a tuned tactic, + captured and replayed). The other cells run the default tactic. + + Not certified: N or K not a multiple of 128 (the kernel needs K % 16 == 0 and N % 8 == 0 + for its TMA alignment), sm_103. + """ + + def __call__( + self, + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype = torch.bfloat16, + use_tvm_ffi: bool = True, + ) -> torch.Tensor: + return torch.ops.trtllm.cute_dsl_fp8_gemm_blackwell( + input, weight, input_scale, weight_scale, output_dtype, use_tvm_ffi + ) + + def reference( + self, + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype = torch.bfloat16, + use_tvm_ffi: bool = True, + ) -> torch.Tensor: + """Both operands dequantized to fp32, one IEEE fp32 matmul, one rounding to + `output_dtype`. `use_tvm_ffi` is ignored.""" + a = dequantize_1x128(input, input_scale) + b = dequantize_128x128(weight, weight_scale) + with true_fp32_matmul(): + out = a @ b.t() + return out.to(output_dtype) + + def is_valid( + self, + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype = torch.bfloat16, + use_tvm_ffi: bool = True, + ) -> None: + """What the kernel assumes and never checks, read from its launch code. + + Every operand reaches the kernel as a bare pointer: a wrong dtype, a strided view, a + misaligned slice or a scale of the wrong shape is read as if it were the layout + above, with no error. The `M == K/128` transposed-scale case is the one this cannot + see. + """ + assert input.dtype == torch.float8_e4m3fn and weight.dtype == torch.float8_e4m3fn, ( + f"input and weight must be float8_e4m3fn, got {input.dtype} and {weight.dtype}; " + "other dtypes are read as e4m3 bytes" + ) + assert input_scale.dtype == torch.float32 and weight_scale.dtype == torch.float32, ( + f"both scales must be fp32 (the checkpoint's weight_scale_inv as stored), got " + f"{input_scale.dtype} and {weight_scale.dtype}; packed UE8M0 scales are misread" + ) + for name, tensor in ( + ("input", input), + ("weight", weight), + ("input_scale", input_scale), + ("weight_scale", weight_scale), + ): + assert tensor.is_contiguous(), ( + f"{name} must be contiguous; a strided view is read as if dense" + ) + assert tensor.data_ptr() % 16 == 0, ( + f"{name} must start 16-byte aligned; the kernel assumes it" + ) + m, k = input.shape + n = weight.shape[0] + assert weight.shape[1] == k, ( + f"input has K={k} but weight has K={weight.shape[1]}; the kernel takes K from " + "weight and reads input with the wrong row length" + ) + assert input_scale.shape == (_ceil_blocks(k), m), ( + f"input_scale must be K-block-major [ceil(K/128), M] = {(_ceil_blocks(k), m)}, got " + f"{tuple(input_scale.shape)}; pass the scales fp8_quantize_1x128 returns" + ) + assert weight_scale.shape == (_ceil_blocks(n), _ceil_blocks(k)), ( + f"weight_scale must be [ceil(N/128), ceil(K/128)] = " + f"{(_ceil_blocks(n), _ceil_blocks(k))}, got {tuple(weight_scale.shape)}" + ) + + +fp8_blockwise_gemm = Fp8BlockwiseGemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/gemm/low_m_bf16_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/gemm/low_m_bf16_gemm.py new file mode 100644 index 000000000000..bac44c6e465f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/gemm/low_m_bf16_gemm.py @@ -0,0 +1,169 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Small-M BF16 projections on FlashInfer's SM10x direct or split-K CuTe DSL GEMM.""" + +from __future__ import annotations + +from typing import Dict, Optional, Tuple + +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 + +# (N, K) -> ((largest M, kernel, tactic), ...) in ascending M; the first row whose M bound +# covers the call wins. Measured per M on B200 against cuBLAS (cold L2, graph replay): every +# row beat cuBLAS by 0.7 to 3.0 us. Shapes and M values outside the table take cuBLAS. +_CHOICES: Dict[Tuple[int, int], Tuple[Tuple[int, str, tuple], ...]] = { + # KDA [f_a | g_a | b] (128 + 128 + 16 local heads), the gate path of every KDA layer. + (272, 4096): ( + (1, "direct", (256, 2, 1)), + (2, "direct", (256, 2, 2)), + (4, "direct", (256, 1, 4)), + (8, "direct", (64, 1, 8)), + (12, "direct", (128, 2, 4)), + (32, "splitk", (64, 8, 4, 11)), + ), + # KDA o_proj, row-parallel: 16 local heads x 128 -> hidden. + (4096, 2048): ( + (2, "direct", (64, 2, 2)), + (4, "splitk", (128, 8, 2, 6)), + (8, "splitk", (64, 8, 1, 10)), + (16, "splitk", (64, 16, 1, 9)), + (32, "splitk", (64, 32, 1, 9)), + ), + # Indexer [k | compress gate | head weights] (128 + 128 + 32), replicated. + (288, 4096): ( + (1, "direct", (256, 2, 1)), + (2, "direct", (256, 2, 2)), + (4, "direct", (128, 1, 4)), + (8, "direct", (64, 1, 8)), + (12, "direct", (128, 2, 4)), + (32, "splitk", (64, 8, 4, 11)), + ), + # MLA [q_a | kv_a] (1536 + 512), replicated, on its dequantized BF16 copy. + (2048, 4096): ( + (1, "direct", (256, 2, 1)), + (2, "direct", (128, 2, 2)), + (8, "splitk", (128, 8, 4, 6)), + (16, "splitk", (64, 16, 2, 10)), + (32, "splitk", (64, 32, 2, 9)), + ), + # MLA o_proj, row-parallel: 16 local heads x 256 -> hidden, dequantized BF16 copy. + (4096, 4096): ( + (1, "direct", (256, 2, 1)), + (8, "splitk", (64, 8, 2, 11)), + (16, "splitk", (64, 16, 2, 11)), + (32, "splitk", (64, 32, 2, 9)), + ), + # LM head, vocabulary-parallel: 154880 / 4 rows. The MTP draft heads run 1 row, verify 4 per request. + (38720, 4096): ( + (1, "direct", (128, 1, 1)), + (2, "direct", (128, 2, 2)), + (8, "splitk", (128, 8, 2, 6)), + (16, "splitk", (64, 16, 1, 11)), + ), + # Indexer wq_b (32 heads x 128) and MLA q_b (16 local heads x 256): both q_lora -> 4096. + (4096, 1536): ( + (8, "splitk", (64, 8, 1, 12)), + (16, "splitk", (64, 16, 1, 11)), + (32, "splitk", (64, 32, 1, 9)), + ), +} + + +def low_m_choice(m: int, n: int, k: int) -> Optional[Tuple[str, tuple]]: + """``(kernel, tactic)`` for an ``[m, k] x [n, k]^T`` product, or None for cuBLAS. + + A direct tactic ``(block_size, outputs_per_block, rows_per_block)`` serves only an M + that its row tile divides; otherwise the next row of the table is tried. + """ + for bound, kernel, tactic in _CHOICES.get((n, k), ()): + if m > bound: + continue + if kernel == "direct" and (tactic[2] > m or m % tactic[2] or n % tactic[1]): + continue + return kernel, tactic + return None + + +class LowMBf16Gemm(OpWrapper): + """``x @ weight.T`` for BF16 ``x [M, K]`` and ``weight [N, K]``, FP32 accumulation rounded once. + + Takes FlashInfer's direct (SIMT, register-prefetched weight) or split-K tensor-core + kernel at the tactic ``low_m_choice`` names, and ``torch.nn.functional.linear`` + (cuBLAS) for any shape or M the table does not cover, so it is safe on every call. + """ + + ARCHS = frozenset({Arch.SM_100}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="KDA [f_a | g_a | b], MTP-3 verify of one request", spec=dict(m=4, n=272, k=4096)), + Cell(why="KDA [f_a | g_a | b], verify of eight requests", spec=dict(m=32, n=272, k=4096)), + Cell(why="KDA o_proj, one request", spec=dict(m=4, n=4096, k=2048)), + Cell(why="KDA o_proj, eight requests", spec=dict(m=32, n=4096, k=2048)), + Cell(why="indexer [k | gate | weights], MTP draft row", spec=dict(m=1, n=288, k=4096)), + Cell(why="indexer [k | gate | weights], one request", spec=dict(m=4, n=288, k=4096)), + Cell(why="indexer wq_b, one request", spec=dict(m=4, n=4096, k=1536)), + Cell(why="indexer wq_b, eight requests", spec=dict(m=32, n=4096, k=1536)), + Cell(why="MLA [q_a | kv_a], one request", spec=dict(m=4, n=2048, k=4096)), + Cell(why="MLA o_proj, one request", spec=dict(m=4, n=4096, k=4096)), + Cell(why="MLA o_proj, MTP draft row", spec=dict(m=1, n=4096, k=4096)), + Cell(why="LM head shard, MTP draft row", spec=dict(m=1, n=38720, k=4096)), + Cell(why="LM head shard, verify of one request", spec=dict(m=4, n=38720, k=4096)), + Cell( + why="LM head shard beyond the table (24 rows): cuBLAS", spec=dict(m=24, n=38720, k=4096) + ), + Cell(why="a shape outside the table: cuBLAS", spec=dict(m=4, n=6144, k=4096)), + ) + + note = """ + The first call of a (kernel, tactic, M) compiles the CuTe DSL kernel, which takes seconds; + the engine's warmup makes those calls before CUDA-graph capture. The output is allocated per + call, which inside a capture is the graph pool's stable address. PDL follows + ``TRTLLM_ENABLE_PDL``. + """ + + def __call__(self, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + # The kernels take DLPack views, which refuse tensors that require grad (parameters). + weight = weight.detach() + m, k = x.shape + n = weight.shape[0] + choice = low_m_choice(m, n, k) if x.is_contiguous() and weight.is_contiguous() else None + if choice is None: + return torch.nn.functional.linear(x, weight) + kernel, tactic = choice + out = torch.empty((m, n), dtype=x.dtype, device=x.device) + if kernel == "direct": + from flashinfer.gemm.kernels.dense_bf16_gemm_direct import ( + DirectTactic, + run_direct_dense, + ) + + return run_direct_dense(x, weight.t(), out, get_env_enable_pdl(), DirectTactic(*tactic)) + from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import ( + SplitKTactic, + run_splitk_dense, + ) + + return run_splitk_dense( + x, weight.t(), None, out, get_env_enable_pdl(), SplitKTactic(*tactic) + ) + + def reference(self, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + with true_fp32_matmul(): + return (x.float() @ weight.float().t()).to(x.dtype) + + def is_valid(self, x: torch.Tensor, weight: torch.Tensor) -> None: + # Non-contiguous operands fall back to cuBLAS rather than reach the kernels, which + # assume row-major x and the stored [N, K] weight. + assert x.dim() == 2 and weight.dim() == 2 and x.shape[1] == weight.shape[1] + assert x.dtype == weight.dtype == torch.bfloat16 + + +low_m_bf16_gemm = LowMBf16Gemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/mhc/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/mhc/__init__.py new file mode 100644 index 000000000000..4fd482cde23b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/mhc/__init__.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The mHC (manifold hyper-connection) boundaries around GLM-5.3-Flash's sublayers, one op per kernel. + +Every sublayer reads and writes four residual streams ``[M, 4, hidden]``. With ``hc`` a sublayer's +mixer weights (``fn [24, 4 * hidden]``, ``scale [3]``, ``base [24]``, all FP32), ``norm`` its input +RMSNorm weight (BF16) and ``scalars`` the tuple ``(sinkhorn_iters, norm_eps, hc_eps, sinkhorn_eps, +post_mult, layer_input_eps)`` (GLM-5.3-Flash's values are ``mhc_big_fuse.GLM_SCALARS``): + +* The first sublayer's pre-mapping, on the embedding's expanded streams:: + + mixes, sqrsum = mhc_gemm_sqrsum(streams.view(M, -1), hc.fn, gemm_sqrsum_tile_n(M)) + bs = big_fuse_block_size(M) + post, comb, x = mhc_big_fuse(mixes, sqrsum, streams, hc.scale, hc.base, norm, *scalars, bs) + +* Every boundary between two sublayers, ``out`` being the previous one's reduced output:: + + args = (hc.fn, hc.scale, hc.base, norm, *scalars) + tactic = fused_hc_tactic(out, streams, post, comb, *args) + scratch = fused_hc_scratch(M, tactic, out.device) + streams, post, comb, x = mhc_fused_hc(out, streams, post, comb, *args, tactic, *scratch) + +* After the last sublayer:: + + streams = mhc_post_mapping(out, streams, post, comb) + +Expanding the embedding into streams and collapsing them at the end (an unweighted mean, then the +final RMSNorm) are plain torch. +""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/mhc/mhc_big_fuse.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/mhc/mhc_big_fuse.py new file mode 100644 index 000000000000..606e19d1a859 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/mhc/mhc_big_fuse.py @@ -0,0 +1,280 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The mHC pre-mapping's coefficients and layer input, from the mixer product (``mhc_big_fuse``). + +The second of the two kernels of an unfused pre-mapping (``mHC.pre_mapping``): +``torch.ops.trtllm.mhc_big_fuse`` (``mhcBigFuseKernel``, one CTA per row). From ``mhc_gemm_sqrsum``'s +raw mixer product ``mixes [M, 24]`` and row sum of squares ``sqrsum [M]`` of the streams +``residual [M, 4, hidden]``, with ``K = 4 * hidden``:: + + rstd = rsqrt(sqrsum[:, None] / K + norm_eps) # the RMS of the flattened streams + pre = sigmoid(mixes[:, 0:4] * rstd * scale[0] + base[0:4]) + hc_eps + post = sigmoid(mixes[:, 4:8] * rstd * scale[1] + base[4:8]) * post_mult # [M, 4] + logits = (mixes[:, 8:24] * rstd * scale[2] + base[8:24]).view(M, 4, 4) + comb = sinkhorn(logits, sinkhorn_iters, sinkhorn_eps) # [M, 4, 4] + layer_input = bf16((pre[:, :, None] * residual).sum(1)) # [M, hidden] + if norm_weight is not None: # the sublayer's input RMSNorm, fused + rms = rsqrt(mean(layer_input**2) + layer_input_eps) + layer_input = bf16(layer_input * rms * norm_weight) + +``comb`` is ``mixes[:, 8:24]`` read row-major as ``[source stream, destination stream]``. +``sinkhorn`` takes a softmax over each row, adds ``sinkhorn_eps`` and normalizes the columns, then +runs ``sinkhorn_iters - 1`` pairs of row and column normalizations; every normalization divides by +the sum plus ``sinkhorn_eps``. It ends on a column normalization, so each destination stream gets a +convex combination of the sources (up to ``sinkhorn_eps``) and the rows sum to one only +approximately. + +Warp 0 computes the coefficients (lanes 0 to 3, one comb row each); the other warps collapse the +streams into the layer input, eight values per thread and step, in FP32 with ``j = 0..3`` as fused +multiply-adds, rounded to BF16. The fused RMSNorm then squares that BF16 layer input, reduces over +the row and rounds the normalized product to BF16 again: the bits a separate RMSNorm of the +unfused pre-mapping's BF16 output would see. The kernels are compiled with fast math (approximate +``exp``, ``rsqrt`` and division). +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: GLM-5.3-Flash's hyper-connection: 4 streams of 4096, 24 mixer outputs. +MULT = 4 +HIDDEN = 4096 +MIX = (2 + MULT) * MULT + +#: The CTA sizes the kernel is instantiated for; 512 at most 64 rows also bulk-copies the row's +#: streams into shared memory first. +BLOCK_SIZE_OPTIONS = (128, 256, 512) + +#: GLM-5.3-Flash's mHC scalars, as the config sets them: ``hc_sinkhorn_iters``; ``rms_norm_eps`` for +#: the mixer's RMS and for the sublayer input RMSNorm; ``hc_eps`` for ``pre`` and the Sinkhorn +#: normalization; ``post`` weights of ``2 * sigmoid``. +GLM_SCALARS = dict( + sinkhorn_iters=20, + norm_eps=1e-5, + hc_eps=1e-6, + sinkhorn_eps=1e-6, + post_mult=2.0, + layer_input_eps=1e-5, +) + + +def big_fuse_block_size(rows: int) -> int: + """The kernel's own default CTA size for ``rows`` rows (``selectBigFuseBlockSize``). + + 128 threads up to 16 rows, 256 above. It is what ``mHC.pre_mapping`` runs when the autotuner + has no entry for the step. + """ + return 128 if rows <= 16 else 256 + + +def sinkhorn(logits: torch.Tensor, iters: int, eps: float) -> torch.Tensor: + """Sinkhorn normalization of ``[..., 4, 4]`` logits as the mHC kernels define it. + + A softmax over the last dim plus ``eps``, a column normalization, then ``iters - 1`` row and + column normalization pairs, each dividing by the sum plus ``eps``. + """ + comb = logits.softmax(-1) + eps + comb = comb / (comb.sum(-2, keepdim=True) + eps) + for _ in range(iters - 1): + comb = comb / (comb.sum(-1, keepdim=True) + eps) + comb = comb / (comb.sum(-2, keepdim=True) + eps) + return comb + + +class MhcBigFuse(OpWrapper): + """A sublayer's mHC coefficients and its (optionally RMS-normalized) layer input. + + Args: + mixes: ``[M, 24]`` FP32, contiguous: ``mhc_gemm_sqrsum``'s mixer product of ``residual``. + sqrsum: ``[M]`` FP32, contiguous: the row sum of squares of the flattened ``residual``. + residual: ``[M, 4, hidden]`` BF16, contiguous: the streams the sublayer reads. + scale: ``[3]`` FP32, contiguous: the sublayer's ``hc_*_scale`` (pre, post, comb). + base: ``[24]`` FP32, contiguous: the sublayer's ``hc_*_base``. + norm_weight: ``[hidden]`` BF16, contiguous, the sublayer's input RMSNorm weight to fuse; or + None to return the collapsed streams unnormalized. + sinkhorn_iters: Sinkhorn iterations, the initial column normalization included (20). + norm_eps: epsilon of the mixer's RMS over the flattened streams (``rms_norm_eps``). + hc_eps: added to ``pre`` (``hc_eps``). + sinkhorn_eps: Sinkhorn epsilon (``hc_eps``). + post_mult: multiplier of the ``post`` sigmoid (2.0). + layer_input_eps: epsilon of the fused RMSNorm (``rms_norm_eps``); ignored without + ``norm_weight``. + block_size: threads per CTA, one of ``BLOCK_SIZE_OPTIONS``; ``big_fuse_block_size(M)`` is + the kernel's default. It changes the layer input's norm reduction order only. + + Returns: + ``(post, comb, layer_input)``: new ``[M, 4]`` and ``[M, 4, 4]`` FP32 and ``[M, hidden]`` + BF16 tensors -- what ``mHC.pre_mapping`` returns (``post`` without its trailing unit + dimension), followed by the input RMSNorm when ``norm_weight`` is given. + """ + + ARCHS = frozenset({Arch.SM_100}) + + _GLM = dict(mult=MULT, hidden=HIDDEN, **GLM_SCALARS) + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode of one request without MTP, the input RMSNorm fused", + spec=dict(rows=1, norm=True, **_GLM), + ), + Cell(why="MTP-3 verify of one request (4 rows)", spec=dict(rows=4, norm=True, **_GLM)), + Cell(why="MTP-3 verify of two requests (8 rows)", spec=dict(rows=8, norm=True, **_GLM)), + Cell( + why="MTP-3 verify of 8 requests (32 rows): the default CTA is 256 threads from 17 rows", + spec=dict(rows=32, norm=True, **_GLM), + ), + Cell( + why="MTP-3 verify of 64 requests (256 rows): no bulk copy, which stops at 64 rows", + spec=dict(rows=256, norm=True, **_GLM), + ), + Cell( + why="a quarter of a 13,312-token prompt: one rank's rows under sequence parallelism", + spec=dict(rows=3328, norm=True, **_GLM), + ), + Cell( + why="MTP-3 verify of one request without the fused norm: mHC.pre_mapping's output", + spec=dict(rows=4, norm=False, **_GLM), + ), + Cell( + why="a quarter of a 13,312-token prompt without the fused norm", + spec=dict(rows=3328, norm=False, **_GLM), + ), + ) + + note = """ + The kernel is compiled for 4 streams. This entry runs it with one split of the mixer product + (`mhc_gemm_sqrsum`'s output); the split-K partial sums the fused boundary feeds it stay inside + `mhc_fused_hc`. The lagged pre-mix variant (`pre_mix_ext` / `pre_mix_out`) is not used. Launched + with programmatic dependent launch when `TRTLLM_ENABLE_PDL` allows it; no rows launch nothing. + `post`, `comb` and, without the fused norm, `layer_input` are the same bits at every block + size. Outputs are allocated per call, from the graph's private pool under CUDA graph capture. + """ + + def __call__( + self, + mixes: torch.Tensor, + sqrsum: torch.Tensor, + residual: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + norm_weight: Optional[torch.Tensor], + sinkhorn_iters: int, + norm_eps: float, + hc_eps: float, + sinkhorn_eps: float, + post_mult: float, + layer_input_eps: float, + block_size: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + rows, hidden = residual.shape[0], residual.shape[-1] + device = residual.device + post = torch.empty((rows, MULT), dtype=torch.float32, device=device) + comb = torch.empty((rows, MULT, MULT), dtype=torch.float32, device=device) + layer_input = torch.empty((rows, hidden), dtype=torch.bfloat16, device=device) + torch.ops.trtllm.mhc_big_fuse( + mixes, + sqrsum, + residual, + scale, + base, + post, + comb, + layer_input, + rows, + MULT * hidden, + hidden, + norm_eps, + hc_eps, + sinkhorn_eps, + post_mult, + sinkhorn_iters, + 1, + block_size, + norm_weight, + layer_input_eps, + None, + None, + ) + return post, comb, layer_input + + def reference( + self, + mixes: torch.Tensor, + sqrsum: torch.Tensor, + residual: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + norm_weight: Optional[torch.Tensor], + sinkhorn_iters: int, + norm_eps: float, + hc_eps: float, + sinkhorn_eps: float, + post_mult: float, + layer_input_eps: float, + block_size: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """The same outputs in FP32 torch. ``block_size`` is ignored. + + Rounded to BF16 only where the unfused pre-mapping stores a tensor: the collapsed layer + input, which a separate RMSNorm reads, and the normalized result. + """ + rows, hidden = residual.shape[0], residual.shape[-1] + rstd = torch.rsqrt(sqrsum.float().reshape(rows, 1) / (MULT * hidden) + norm_eps) + scaled = mixes.float().reshape(rows, MIX) * rstd + s, b = scale.float().reshape(-1), base.float().reshape(-1) + pre = torch.sigmoid(scaled[:, :MULT] * s[0] + b[:MULT]) + hc_eps + post = torch.sigmoid(scaled[:, MULT : 2 * MULT] * s[1] + b[MULT : 2 * MULT]) * post_mult + logits = (scaled[:, 2 * MULT :] * s[2] + b[2 * MULT : MIX]).reshape(rows, MULT, MULT) + comb = sinkhorn(logits, sinkhorn_iters, sinkhorn_eps) + streams = residual.float().reshape(rows, MULT, hidden) + collapsed = pre[:, :1] * streams[:, 0] + for j in range(1, MULT): + collapsed += pre[:, j : j + 1] * streams[:, j] + layer_input = collapsed.to(torch.bfloat16) + if norm_weight is not None: + li = layer_input.float() + li = li * torch.rsqrt(li.square().mean(-1, keepdim=True) + layer_input_eps) + layer_input = (li * norm_weight.float()).to(torch.bfloat16) + return post, comb, layer_input + + def is_valid( + self, + mixes: torch.Tensor, + sqrsum: torch.Tensor, + residual: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + norm_weight: Optional[torch.Tensor], + sinkhorn_iters: int, + norm_eps: float, + hc_eps: float, + sinkhorn_eps: float, + post_mult: float, + layer_input_eps: float, + block_size: int, + ) -> None: + assert residual.dim() == 3 and residual.shape[1] == MULT, ( + f"the kernel is compiled for {MULT} streams: residual is [M, {MULT}, hidden]; " + f"got {tuple(residual.shape)}" + ) + rows, hidden = residual.shape[0], residual.shape[2] + assert hidden % 8 == 0, "each thread moves 8 BF16 values per step; the kernel has no tail" + assert mixes.shape == (rows, MIX) and sqrsum.shape == (rows,), ( + f"one split: mixes [{rows}, {MIX}] and sqrsum [{rows}]; " + f"got {tuple(mixes.shape)} and {tuple(sqrsum.shape)}" + ) + assert scale.numel() >= 3 and base.numel() >= MIX, "scale holds 3 values and base 24" + for name, t in ( + ("mixes", mixes), + ("sqrsum", sqrsum), + ("residual", residual), + ("scale", scale), + ("base", base), + ): + assert t.is_contiguous(), f"{name} is read with dense row-major strides" + + +mhc_big_fuse = MhcBigFuse() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/mhc/mhc_fused_hc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/mhc/mhc_fused_hc.py new file mode 100644 index 000000000000..447ac38df007 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/mhc/mhc_fused_hc.py @@ -0,0 +1,672 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The boundary between two mHC-wrapped sublayers in one call (``mhc_fused_hc``). + +``torch.ops.trtllm.mhc_fused_hc``, the op ``mHC.fused_hc`` calls: the previous sublayer's +post-mapping into the residual streams, then the next sublayer's pre-mapping of the new streams, +with that sublayer's input RMSNorm folded in. For ``M`` rows of ``mult = 4`` streams of ``hidden`` +values it computes ``mhc_post_mapping`` followed by ``mhc_gemm_sqrsum`` and ``mhc_big_fuse`` on the +BF16 streams it stores (each of those modules states its step's formula):: + + streams = bf16(post[:, :, None] * x[:, None, :] + comb.mT @ residual) # returned + mixes = streams.view(M, 4 * hidden) @ fn.T # [M, 24] + sqrsum = streams.view(M, 4 * hidden).square().sum(-1) # [M] + post_out, comb_out, layer_input = mhc_big_fuse(mixes, sqrsum, streams, ...) + +The op runs one of four backends, picked by ``tactic``: + +* ``fused_half_mma``: two kernels. A tcgen05 TF32 GEMM fused with the post-mapping writes the + streams and, per hidden split, a partial mixer product and sum of squares of the BF16-rounded + streams; ``mhcBigFuseKernel`` sums the splits in a fixed order and finishes the boundary. +* ``fused_half_fma``: two kernels. The same on CUDA-core FP32 FMA, over the FP32 streams before + their BF16 rounding, with ``hidden`` split into at most 8 slices. +* ``fused_all_mma``: one kernel (plus a counter-zeroing kernel when split): the TF32 GEMM, a + barrier per 64 rows, then the split sum and the rest of the boundary inline. +* ``fused_all_fma``: one kernel (plus a workspace-zeroing kernel): FP32 FMA, whose at most two + hidden splits are added onto one zeroed slot -- two addends sum the same in either order -- and + whose last CTA of each row group finishes the boundary. + +The TF32 backends read ``fn`` at TF32 precision. The checkpoint's mixer weights are BF16 values, +which TF32 holds exactly, as it does the BF16 streams. Every tactic gives the same bits run to +run; tactics sum in different orders (and the FMA backends give the mixer the unrounded streams), +so they differ from each other in the last bits. + +``tactic`` is ``(backend, tile_n, num_k_splits, bigfuse_block_size, tile_m)``. ``fused_hc_tactic`` +picks it as ``mHC.fused_hc`` does: the autotuner over ``fused_hc_tactics`` (one tactic for every TP +rank), the tensor-core fallback on a cache miss, and ``DETERMINISTIC_TACTIC`` under +``FORCE_DETERMINISTIC=1`` or ``FORCE_MHC_KERNEL_DETERMINISTIC=1``. ``fused_hc_scratch`` allocates +the three workspaces a tactic needs. +""" + +from __future__ import annotations + +import math +import os +from functools import lru_cache +from typing import Optional, Sequence + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.autotuner import ( + AutoTuner, + ConstraintSpec, + DistributedTuningStrategy, + DynamicTensorSpec, + OptimizationProfile, + TunableRunner, + TuningConfig, +) +from tensorrt_llm._utils import get_sm_version + +from .mhc_big_fuse import GLM_SCALARS, mhc_big_fuse +from .mhc_gemm_sqrsum import mhc_gemm_sqrsum +from .mhc_post_mapping import mhc_post_mapping + +#: GLM-5.3-Flash's hyper-connection: 4 streams of 4096, 24 mixer outputs. +MULT = 4 +HIDDEN = 4096 +MIX = (2 + MULT) * MULT + +#: ``tactic[0]`` -> the op's backend code. +BACKENDS = {"fused_half_mma": 0, "fused_half_fma": 1, "fused_all_mma": 2, "fused_all_fma": 3} + +#: The tactic under deterministic mode (``deterministic_requested``) and the fallback without the +#: tensor-core backends: +#: half-fused FMA, two columns per CTA, no hidden split, 256-thread BigFuse. It has no cross-CTA +#: floating-point atomics and one reduction order for every row count. +DETERMINISTIC_TACTIC = ("fused_half_fma", 2, 1, 256, 1) + +#: The autotuner's key for this op's tactic choice. +TUNING_OP = "glm5_next::mhc_fused_hc" + +#: A tactic: ``(backend, tile_n, num_k_splits, bigfuse_block_size, tile_m)``. +Tactic = tuple[str, int, int, int, int] + +#: The op's arguments besides the tensors the autotuner sees: the RMSNorm weight and the scalars. +_RunnerArg = torch.Tensor | int | float | None + +# (tile_n, num_k_splits) pairs of the half-fused FMA kernel (C++ pickFhcFma). +_HALF_FMA_TN_KS = ( + (1, 1), + (1, 2), + (1, 4), + (1, 8), + (2, 1), + (2, 2), + (2, 4), + (2, 8), + (3, 1), + (3, 2), + (3, 4), + (4, 1), + (4, 2), + (6, 1), + (8, 1), + (12, 1), + (24, 1), +) +# (tile_n, num_k_splits, tile_m) triples of the all-in-one FMA kernel (C++ pickFhcFmaAllInOne). +_ALL_FMA_TN_KS_TM = tuple( + (tn, ks, tm) + for tm in (1, 2, 4) + for tn, ks in ((1, 1), (1, 2), (2, 1), (2, 2), (3, 1), (4, 1), (6, 1), (8, 1), (12, 1), (24, 1)) +) +# Split counts the TF32 kernels are instantiated for; a hidden size takes those that divide its +# 64-wide tiles evenly. +_MMA_KS = (1, 2, 4, 5, 7, 8, 10, 14, 16, 20, 28, 32, 40, 56, 64, 80, 112) +_MMA_HIDDEN_SIZES = (4096, 5120, 7168) +_MMA_BLOCK_M = 64 +_BIGFUSE_BLOCK_SIZES = (128, 256, 512) +# The candidate filters cap a grid in units of B200's 148 SMs, on every SKU. +_GRID_UNIT = 148 + + +@lru_cache(maxsize=1) +def fused_hc_mma_enabled() -> bool: + """Whether the tensor-core backends run here: an SM 10.x GPU and a build with DeepGEMM.""" + return 100 <= get_sm_version() < 110 and bool(torch.ops.trtllm.mhc_fused_hc_mma_enabled()) + + +def _mma_split_options(hidden: int) -> tuple[int, ...]: + if hidden not in _MMA_HIDDEN_SIZES: + return () + tiles = hidden // 64 + return tuple(ks for ks in _MMA_KS if ks <= tiles and tiles % ks == 0) + + +def _mma_split_ladder(rows: int) -> int: + """The half-fused TF32 kernel's default split count (C++ ``mhcFusedHcMmaKSplits``).""" + for max_rows, splits in ((512, 16), (1024, 8), (2048, 4), (4096, 2)): + if rows <= max_rows: + return splits + return 1 + + +def _mma_bigfuse_block_size(rows: int) -> int: + """The half-fused TF32 kernel's default BigFuse CTA size (C++ ``selectBigFuseBS``).""" + if rows >= 4096: + return 128 + return 256 if rows >= 256 else 512 + + +def deterministic_requested() -> bool: + """Whether the environment asks for run-to-run identical mHC results, as ``mHC.fused_hc`` + reads it: ``FORCE_DETERMINISTIC=1`` or ``FORCE_MHC_KERNEL_DETERMINISTIC=1``.""" + return ( + os.environ.get("FORCE_DETERMINISTIC", "0") == "1" + or os.environ.get("FORCE_MHC_KERNEL_DETERMINISTIC", "0") == "1" + ) + + +def default_fused_hc_tactic(rows: int, hidden: int, mma_enabled: bool) -> Tactic: + """The tactic of a call the autotuner has no entry for: ``mHC.fused_hc``'s fallback, resolved. + + The half-fused TF32 kernel at its own default split count and BigFuse CTA size where the + tensor-core backends run, ``DETERMINISTIC_TACTIC`` elsewhere. + """ + if mma_enabled and hidden in _MMA_HIDDEN_SIZES: + return ("fused_half_mma", 0, _mma_split_ladder(rows), _mma_bigfuse_block_size(rows), 1) + return DETERMINISTIC_TACTIC + + +def fused_hc_tactics( + rows: int, hidden: int, mma_enabled: bool, pin_small_m_half_fma: bool +) -> list[Tactic]: + """The tactics ``mHC.fused_hc``'s autotuner profiles for ``rows`` rows (a tuning bucket). + + Up to 32 rows: the all-in-one FMA kernel at every tile shape whose grid stays within four + waves of 148 CTAs, and the half-fused FMA kernel at every tile width, split count and BigFuse + CTA size (split ones only while ``rows * 24 / tile_n`` is under 296). From 32 rows, with + tensor cores: the half-fused TF32 kernel at the largest split count of at most + ``128 / ceil(rows / 64)`` and its measured BigFuse CTA sizes, and from 64 rows the all-in-one + TF32 kernel at that split count. Above 32 rows without them: the half-fused FMA kernel unsplit. + + ``pin_small_m_half_fma`` (SM 10.3) keeps the small-row half-fused FMA candidates to one split + count (4 up to 4 rows, else 2) and one BigFuse CTA size (512), so that those candidates differ + only in tile width, which assigns columns to CTAs without changing any sum. + """ + tactics: dict[Tactic, None] = {} + mma_splits = _mma_split_options(hidden) + mma_ok = mma_enabled and bool(mma_splits) + if rows <= 32: + for tn, ks, tm in _ALL_FMA_TN_KS_TM: + if -(-rows // tm) * (MIX // tn) * ks <= _GRID_UNIT * 4: + tactics[("fused_all_fma", tn, ks, 0, tm)] = None + pinned_ks = (4 if rows <= 4 else 2) if pin_small_m_half_fma else None + block_sizes = (512,) if pin_small_m_half_fma else _BIGFUSE_BLOCK_SIZES + for tn, ks in _HALF_FMA_TN_KS: + if pinned_ks is not None and ks != pinned_ks: + continue + if ks > 1 and rows * (MIX // tn) >= _GRID_UNIT * 2: + continue + for bs in block_sizes: + tactics[("fused_half_fma", tn, ks, bs, 1)] = None + if mma_ok and rows >= 32: + m_tiles = -(-rows // _MMA_BLOCK_M) + fitting = [ks for ks in mma_splits if ks <= max(1, 128 // m_tiles)] + ks = fitting[-1] if fitting else mma_splits[0] + if m_tiles * ks <= _GRID_UNIT * 4: + if rows < 512: + block_sizes = (512,) + elif rows == 1024: + block_sizes = (128, 512) + elif rows >= 8192: + block_sizes = (128, 256) + else: + block_sizes = (256,) + for bs in block_sizes: + tactics[("fused_half_mma", 0, ks, bs, 1)] = None + if rows >= 64: + tactics[("fused_all_mma", 0, ks, 0, 1)] = None + if not mma_ok and rows > 32: + bs = 512 if rows < 512 else 256 + for tn, ks in _HALF_FMA_TN_KS: + if ks == 1: + tactics[("fused_half_fma", tn, ks, bs, 1)] = None + return list(tactics) + + +def _splits(rows: int, tactic: Sequence) -> int: + backend, _, num_k_splits, _, _ = tactic + if backend == "fused_half_mma": + return num_k_splits if num_k_splits > 0 else _mma_split_ladder(rows) + if backend in ("fused_all_mma", "fused_half_fma"): + return max(num_k_splits, 1) + return 1 # fused_all_fma adds its at most two splits onto one slot + + +def fused_hc_scratch_shapes( + rows: int, tactic: Sequence +) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + """Shapes of ``(y_acc_ws, r_acc_ws, done_counter_ws)`` for ``rows`` rows under ``tactic``. + + One FP32 ``[M, 24]`` mixer partial and ``[M]`` sum-of-squares partial per hidden split + (``[splits, M, 24]`` and ``[splits, M]`` when split; the all-in-one FMA kernel always keeps one + slot), and int32 done counters: one per 64 rows for ``fused_all_mma``, one per ``tile_m`` rows + for ``fused_all_fma``, one unused counter for the half-fused backends. + """ + backend, _, _, _, tile_m = tactic + splits = _splits(rows, tactic) + if backend == "fused_all_mma": + counters = -(-rows // _MMA_BLOCK_M) + elif backend == "fused_all_fma": + counters = -(-rows // max(1, tile_m)) + else: + counters = 1 + if splits > 1: + return (splits, rows, MIX), (splits, rows), (counters,) + return (rows, MIX), (rows,), (counters,) + + +def fused_hc_scratch( + rows: int, tactic: Sequence, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """New ``(y_acc_ws, r_acc_ws, done_counter_ws)`` for ``rows`` rows under ``tactic``. + + Uninitialized: the kernels write every slot they read, and the all-in-one backends zero what + they accumulate into. + """ + y_shape, r_shape, done_shape = fused_hc_scratch_shapes(rows, tactic) + return ( + torch.empty(y_shape, dtype=torch.float32, device=device), + torch.empty(r_shape, dtype=torch.float32, device=device), + torch.empty(done_shape, dtype=torch.int32, device=device), + ) + + +def _pins_small_m_half_fma(device: torch.device) -> bool: + return torch.cuda.get_device_capability(device) == (10, 3) + + +def _tuning_buckets(rows: int) -> tuple[int, ...]: + buckets = (1, 2, 4, 8, 16, 32, 64, 128) + if rows >= 128: + buckets += tuple(range(256, max(1024, min(rows, 8192)) + 1, 128)) + return buckets + + +def _tuning_bucket(rows: int) -> int: + if rows <= 128: + bucket = 1 + while bucket < rows: + bucket *= 2 + return bucket + return -(-rows // 128) * 128 + + +def _rows_of_x(shapes: Sequence[Sequence[int]]) -> int: + return shapes[0][0] + + +class _FusedHcRunner(TunableRunner): + """The autotuner's view of ``mhc_fused_hc``: the candidates of a row bucket, and one call each. + + The inputs are ``[x, residual, post, comb, fn, scale, base]``; the rest of the op's arguments + arrive as keyword arguments. Rows are bucketed by powers of two up to 128, then in steps of 128. + + The TP ranks share the profiling (``PARALLEL``, as ``mHC.fused_hc``'s runner does): each rank + times its share of the candidates and every rank keeps the fastest of all, so every rank runs + the same tactic. In decode every rank holds the same rows, and a tactic sums in its own order, + so ranks on different tactics would carry streams that differ in the last bits and could route + a row to different experts. The ranks call the autotuner in lockstep: the target pads a + sequence-parallel step's rows to the same count on every rank. + """ + + tuning_config = TuningConfig( + dynamic_tensor_specs=( + DynamicTensorSpec( + input_idx=0, + dim_idx=0, + gen_tuning_buckets=_tuning_buckets, + map_to_tuning_buckets=_tuning_bucket, + ), + ), + constraint_specs=tuple( + ConstraintSpec(input_idx=i, dim_idx=0, infer_shape=_rows_of_x) for i in (1, 2, 3) + ), + distributed_tuning_strategy=DistributedTuningStrategy.PARALLEL, + ) + + def __init__(self, hidden: int) -> None: + self.hidden = hidden + + def unique_id(self) -> tuple[int, int]: + return (MULT, self.hidden) + + def get_valid_tactics( + self, inputs: list[torch.Tensor], profile: OptimizationProfile, **kwargs: _RunnerArg + ) -> list[Tactic]: + x = inputs[0] + return fused_hc_tactics( + x.shape[0], self.hidden, fused_hc_mma_enabled(), _pins_small_m_half_fma(x.device) + ) + + def forward( + self, inputs: list[torch.Tensor], *, tactic: Sequence | int = -1, **kwargs: _RunnerArg + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + x = inputs[0] + rows = x.shape[0] + if tactic == -1: + tactic = default_fused_hc_tactic(rows, self.hidden, fused_hc_mma_enabled()) + return mhc_fused_hc( + *inputs, + kwargs["norm_weight"], + kwargs["sinkhorn_iters"], + kwargs["norm_eps"], + kwargs["hc_eps"], + kwargs["sinkhorn_eps"], + kwargs["post_mult"], + kwargs["layer_input_eps"], + tactic, + *fused_hc_scratch(rows, tactic, x.device), + ) + + +_RUNNERS: dict[int, _FusedHcRunner] = {} + + +def fused_hc_tactic( + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + norm_weight: Optional[torch.Tensor], + sinkhorn_iters: int, + norm_eps: float, + hc_eps: float, + sinkhorn_eps: float, + post_mult: float, + layer_input_eps: float, +) -> Tactic: + """The tactic for an ``mhc_fused_hc`` call with these arguments, chosen as ``mHC.fused_hc`` does. + + ``DETERMINISTIC_TACTIC`` when ``deterministic_requested()``. Otherwise the autotuner's choice + under ``TUNING_OP``. In tuning mode it profiles ``fused_hc_tactics`` at every row bucket -- + powers of two up to 128 and, from a call of 128 rows, multiples of 128 up to ``max(1024, + min(M, 8192))`` -- running the op on these weights and on bucket-sized activations, with the + candidates shared out over the TP ranks and the fastest kept on all of them; every rank must + make the call, with the same row count. Outside it, it looks up the call's row bucket (rows + rounded up to a power of two, or above 128 to a multiple of 128) and returns + ``default_fused_hc_tactic`` on a miss. The arguments are ``mhc_fused_hc``'s without the + tactic and the workspaces. + """ + if deterministic_requested(): + return DETERMINISTIC_TACTIC + rows, hidden = x.shape + runner = _RUNNERS.get(hidden) + if runner is None: + runner = _RUNNERS[hidden] = _FusedHcRunner(hidden) + _, tactic = AutoTuner.get().choose_one( + TUNING_OP, + [runner], + _FusedHcRunner.tuning_config, + [x, residual, post, comb, fn, scale, base], + norm_weight=norm_weight, + sinkhorn_iters=sinkhorn_iters, + norm_eps=norm_eps, + hc_eps=hc_eps, + sinkhorn_eps=sinkhorn_eps, + post_mult=post_mult, + layer_input_eps=layer_input_eps, + ) + if tactic == -1: + return default_fused_hc_tactic(rows, hidden, fused_hc_mma_enabled()) + return tuple(tactic) + + +class MhcFusedHc(OpWrapper): + """The previous sublayer's post-mapping, then the next sublayer's pre-mapping and input RMSNorm. + + Args: + x: ``[M, hidden]`` BF16, contiguous: the previous sublayer's output, after its TP reduction. + residual: ``[M, 4, hidden]`` BF16, contiguous: the streams the previous sublayer read. + post: ``[M, 4]`` or ``[M, 4, 1]`` FP32, contiguous: the previous sublayer's post coefficients. + comb: ``[M, 4, 4]`` FP32, contiguous: the previous sublayer's Sinkhorn-normalized combination, + ``[source stream, destination stream]``. + fn: ``[24, 4 * hidden]`` FP32, contiguous: the next sublayer's mixer weights. + scale: ``[3]`` FP32, contiguous: the next sublayer's mixer scales (pre, post, comb). + base: ``[24]`` FP32, contiguous: the next sublayer's mixer biases. + norm_weight: ``[hidden]`` BF16, contiguous: the next sublayer's input RMSNorm weight; or + None for the unnormalized layer input. + sinkhorn_iters: Sinkhorn iterations, the initial column normalization included (20). + norm_eps: epsilon of the mixer's RMS over the flattened streams (``rms_norm_eps``). + hc_eps: added to ``pre`` (``hc_eps``). + sinkhorn_eps: Sinkhorn epsilon (``hc_eps``). + post_mult: multiplier of the ``post`` sigmoid (2.0). + layer_input_eps: epsilon of the fused input RMSNorm (``rms_norm_eps``). + tactic: ``(backend, tile_n, num_k_splits, bigfuse_block_size, tile_m)``, e.g. from + ``fused_hc_tactic``; see ``fused_hc_tactics`` for the candidates. + y_acc_ws: FP32 mixer partials, ``fused_hc_scratch_shapes(M, tactic)[0]`` or larger. + r_acc_ws: FP32 sum-of-squares partials, ``fused_hc_scratch_shapes(M, tactic)[1]`` or larger. + done_counter_ws: int32 counters, ``fused_hc_scratch_shapes(M, tactic)[2]`` or larger. + + Returns: + ``(residual_out, post_out, comb_out, layer_input)``: the new streams (BF16, ``residual``'s + shape), the next sublayer's ``post`` ``[M, 4]`` and ``comb`` ``[M, 4, 4]`` (FP32), and its + layer input ``[M, hidden]`` (BF16, RMS-normalized when ``norm_weight`` is given) -- what + ``mHC.fused_hc`` returns, ``post`` without its trailing unit dimension. All new tensors. + """ + + ARCHS = frozenset({Arch.SM_100}) + + _GLM = dict(mult=MULT, hidden=HIDDEN, **GLM_SCALARS) + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode of one request without MTP: FMA candidates only, one CTA row", + spec=dict(rows=1, norm=True, **_GLM), + ), + Cell(why="MTP-3 verify of one request (4 rows)", spec=dict(rows=4, norm=True, **_GLM)), + Cell(why="MTP-3 verify of two requests (8 rows)", spec=dict(rows=8, norm=True, **_GLM)), + Cell( + why="MTP-3 verify of 8 requests (32 rows): both FMA and TF32 candidates", + spec=dict(rows=32, norm=True, **_GLM), + ), + Cell( + why="MTP-3 verify of 64 requests (256 rows): TF32 candidates at 32 hidden splits", + spec=dict(rows=256, norm=True, **_GLM), + ), + Cell( + why="a quarter of a 13,312-token prompt (sequence parallel): TF32 at 2 hidden splits", + spec=dict(rows=3328, norm=True, **_GLM), + ), + Cell( + why="MTP-3 verify of one request without the fused norm: the unnormalized variants", + spec=dict(rows=4, norm=False, **_GLM), + ), + ) + + note = """ + The cells are driven with every tactic `fused_hc_tactic` can return at their row count: the + autotuner's candidates, the cache-miss fallback and `DETERMINISTIC_TACTIC`. The tensor-core + backends need an SM 10.x GPU, a build with DeepGEMM (`fused_hc_mma_enabled`) and hidden 4096, + 5120 or 7168; the FMA backends take any hidden divisible by 64. The kernels are compiled for 4 + streams. The lagged pre-mix variant (`pre_mix_ext` / `pre_mix_out`) is not used. + + `fused_hc_tactic` tunes as `mHC.fused_hc` does: the TP ranks share the profiling and keep one + tactic, so their streams stay bitwise equal. `FORCE_DETERMINISTIC=1` or + `FORCE_MHC_KERNEL_DETERMINISTIC=1` puts every rank on `DETERMINISTIC_TACTIC`, whose sums do not + depend on the row count either. The candidates of a row count do not depend on the weights or + the scalars, and a step's boundaries share their shapes, so a tactic can be chosen once per + step. + + The op checks the split workspaces of the first three backends against the split count it + runs, and neither the all-in-one FMA kernel's workspaces nor any counter array; `is_valid` + covers those. Nothing needs zeroing beforehand, so successive calls on one stream can share + the workspaces; concurrent calls cannot. The FMA kernels and BigFuse launch with programmatic + dependent launch when `TRTLLM_ENABLE_PDL` allows it. A call of no rows launches nothing. + Outputs are allocated per call, from the graph's private pool under CUDA graph capture; the + tensor-core backends cache their TMA descriptors by address and shape. + """ + + def __call__( + self, + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + norm_weight: Optional[torch.Tensor], + sinkhorn_iters: int, + norm_eps: float, + hc_eps: float, + sinkhorn_eps: float, + post_mult: float, + layer_input_eps: float, + tactic: Sequence, + y_acc_ws: torch.Tensor, + r_acc_ws: torch.Tensor, + done_counter_ws: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + rows, hidden = x.shape + backend, tile_n, num_k_splits, bigfuse_block_size, tile_m = tactic + residual_out = torch.empty_like(residual, memory_format=torch.contiguous_format) + post_out = torch.empty((rows, MULT), dtype=torch.float32, device=x.device) + comb_out = torch.empty((rows, MULT, MULT), dtype=torch.float32, device=x.device) + layer_input = torch.empty((rows, hidden), dtype=torch.bfloat16, device=x.device) + torch.ops.trtllm.mhc_fused_hc( + x, + residual, + post, + comb, + fn, + scale, + base, + residual_out, + post_out, + comb_out, + layer_input, + y_acc_ws, + r_acc_ws, + done_counter_ws, + rows, + hidden, + MULT, + norm_eps, + hc_eps, + sinkhorn_eps, + post_mult, + sinkhorn_iters, + BACKENDS[backend], + tile_n, + num_k_splits, + bigfuse_block_size, + tile_m, + norm_weight, + layer_input_eps, + None, + None, + ) + return residual_out, post_out, comb_out, layer_input + + def reference( + self, + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + norm_weight: Optional[torch.Tensor], + sinkhorn_iters: int, + norm_eps: float, + hc_eps: float, + sinkhorn_eps: float, + post_mult: float, + layer_input_eps: float, + tactic: Sequence, + y_acc_ws: torch.Tensor, + r_acc_ws: torch.Tensor, + done_counter_ws: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """The unfused boundary from the three entries' references; the tactic and workspaces are + ignored. + + ``mhc_post_mapping``'s streams, rounded to BF16 as the next sublayer reads them, then + ``mhc_gemm_sqrsum`` (a true-FP32 mixer product) and ``mhc_big_fuse`` on those streams. + """ + rows = x.shape[0] + streams = mhc_post_mapping.reference(x, residual, post, comb) + mixes, sqrsum = mhc_gemm_sqrsum.reference(streams.reshape(rows, -1), fn, 0) + post_out, comb_out, layer_input = mhc_big_fuse.reference( + mixes, + sqrsum, + streams.reshape(rows, MULT, -1), + scale, + base, + norm_weight, + sinkhorn_iters, + norm_eps, + hc_eps, + sinkhorn_eps, + post_mult, + layer_input_eps, + 0, + ) + return streams, post_out, comb_out, layer_input + + def is_valid( + self, + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + norm_weight: Optional[torch.Tensor], + sinkhorn_iters: int, + norm_eps: float, + hc_eps: float, + sinkhorn_eps: float, + post_mult: float, + layer_input_eps: float, + tactic: Sequence, + y_acc_ws: torch.Tensor, + r_acc_ws: torch.Tensor, + done_counter_ws: torch.Tensor, + ) -> None: + assert x.dim() == 2, f"x is [M, hidden]; got {tuple(x.shape)}" + rows, hidden = x.shape + assert residual.numel() == rows * MULT * hidden and residual.shape[-1] == hidden, ( + f"residual must hold [{rows}, {MULT}, {hidden}]; got {tuple(residual.shape)}" + ) + assert post.numel() == rows * MULT and comb.numel() == rows * MULT * MULT, ( + f"post holds [{rows}, {MULT}] and comb [{rows}, {MULT}, {MULT}]" + ) + assert fn.shape == (MIX, MULT * hidden), ( + f"fn is [{MIX}, {MULT * hidden}]; got {tuple(fn.shape)}" + ) + assert scale.numel() >= 3 and base.numel() >= MIX, "scale holds 3 values and base 24" + for name, t in ( + ("x", x), + ("residual", residual), + ("post", post), + ("comb", comb), + ("fn", fn), + ("scale", scale), + ("base", base), + ("y_acc_ws", y_acc_ws), + ("r_acc_ws", r_acc_ws), + ("done_counter_ws", done_counter_ws), + ): + assert t.is_contiguous(), f"{name} is read with dense row-major strides" + y_shape, r_shape, done_shape = fused_hc_scratch_shapes(rows, tactic) + y_need, r_need = math.prod(y_shape), math.prod(r_shape) + assert y_acc_ws.numel() >= y_need and r_acc_ws.numel() >= r_need, ( + f"{tuple(tactic)} at {rows} rows writes {y_need} mixer and {r_need} sum-of-squares " + f"partials; got {y_acc_ws.numel()} and {r_acc_ws.numel()}" + ) + assert done_counter_ws.numel() >= done_shape[0], ( + f"{tuple(tactic)} at {rows} rows counts in {done_shape[0]} slots; " + f"got {done_counter_ws.numel()}" + ) + + +mhc_fused_hc = MhcFusedHc() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/mhc/mhc_gemm_sqrsum.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/mhc/mhc_gemm_sqrsum.py new file mode 100644 index 000000000000..9003faaeec39 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/mhc/mhc_gemm_sqrsum.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The mHC pre-mapping's mixer product and row sum of squares, on CUDA-core FP32 FMA. + +The first of the two kernels of an unfused pre-mapping (``mHC.pre_mapping``): +``torch.ops.trtllm.mhc_gemm_sqrsum_fma`` (``mhcGemmSqrsumFmaKernel``). For the ``M`` rows of the +residual streams flattened stream-major, ``x = residual.view(M, K)`` with ``K = mult * hidden`` +(BF16), and the next sublayer's mixer ``fn`` (FP32 ``[mix, K]``, ``mix = (2 + mult) * mult = 24``):: + + mixes[t, n] = sum_k x[t, k] * fn[n, k] # [M, 24] FP32, before the RMS scaling + sqrsum[t] = sum_k x[t, k] ** 2 # [M] FP32 + +``mhc_big_fuse`` turns the two into the coefficients and the layer input. One CTA of 256 threads +computes ``tile_n`` columns of one row over all of ``K``: each thread runs a sequential FMA chain over +its 64 elements (4 of every 1024), then a warp butterfly and a sum over the 8 warps in warp order. +Each column is computed the same way whatever the tile width, so ``tile_n`` changes the grid and not +the result. A BF16-by-FP32 product is exact in FP32, so the result differs from any other FP32 +summation of the same products only in rounding order. +""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +#: GLM-5.3-Flash's hyper-connection: 4 streams of 4096, 24 mixer outputs. +MULT = 4 +HIDDEN = 4096 +MIX = (2 + MULT) * MULT + +#: The tile widths the kernel is instantiated for: the divisors of 24. +TILE_N_OPTIONS = (1, 2, 3, 4, 6, 8, 12, 24) + + +def gemm_sqrsum_tile_n(rows: int) -> int: + """The kernel's own default tile width for ``rows`` rows (``selectFmaTileN`` at 24 columns). + + One column per CTA up to 32 rows, eight above. It is what ``mHC.pre_mapping`` runs when the + autotuner has no entry for the step. + """ + return 1 if rows <= 32 else 8 + + +class MhcGemmSqrsum(OpWrapper): + """``(x @ fn.T, (x * x).sum(-1))`` in FP32 for BF16 ``x [M, K]`` and FP32 ``fn [24, K]``. + + Args: + x: ``[M, K]`` BF16, contiguous: the residual streams ``[M, 4, hidden]`` viewed as + ``[M, 4 * hidden]``. + fn: ``[24, K]`` FP32, contiguous: the next sublayer's mixer weights as stored + (``hc_attn_fn`` / ``hc_ffn_fn``; the checkpoint's BF16 values held as FP32). + tile_n: output columns per CTA, one of ``TILE_N_OPTIONS``; ``gemm_sqrsum_tile_n(M)`` is the + kernel's default. It does not change the result. + + Returns: + ``(mixes, sqrsum)``: new ``[M, 24]`` and ``[M]`` FP32 tensors. + """ + + ARCHS = frozenset({Arch.SM_100}) + + _GLM = dict(mult=MULT, hidden=HIDDEN) + CELLS: tuple[Cell, ...] = ( + Cell(why="decode of one request without MTP", spec=dict(rows=1, **_GLM)), + Cell(why="MTP-3 verify of one request (4 rows)", spec=dict(rows=4, **_GLM)), + Cell(why="MTP-3 verify of two requests (8 rows)", spec=dict(rows=8, **_GLM)), + Cell( + why="MTP-3 verify of 8 requests (32 rows): the default tile's last one-column step", + spec=dict(rows=32, **_GLM), + ), + Cell(why="MTP-3 verify of 64 requests (256 rows)", spec=dict(rows=256, **_GLM)), + Cell( + why="a quarter of a 13,312-token prompt: one rank's rows under sequence parallelism", + spec=dict(rows=3328, **_GLM), + ), + ) + + note = """ + The target runs the pre-mapping once per forward step, on the embedding's expanded streams; + every later boundary goes through `mhc_fused_hc`. Only the CUDA-core kernel is wrapped: the + DeepGEMM TF32 product that `mHC.pre_mapping`'s autotuner can also pick is not offered. Each CTA + reads its row of `x` and `tile_n` rows of `fn` in full, so every row of `x` is read + `24 / tile_n` times. Launched without programmatic dependent launch; a call of no rows launches + nothing. Outputs are allocated per call, from the graph's private pool under CUDA graph capture. + """ + + def __call__( + self, x: torch.Tensor, fn: torch.Tensor, tile_n: int + ) -> tuple[torch.Tensor, torch.Tensor]: + rows, k = x.shape + mix = fn.shape[0] + mixes = torch.empty((rows, mix), dtype=torch.float32, device=x.device) + sqrsum = torch.empty((rows,), dtype=torch.float32, device=x.device) + # tile_m is reserved by the launcher and ignored. + torch.ops.trtllm.mhc_gemm_sqrsum_fma(x, fn, mixes, sqrsum, rows, mix, k, tile_n, 0) + return mixes, sqrsum + + def reference( + self, x: torch.Tensor, fn: torch.Tensor, tile_n: int + ) -> tuple[torch.Tensor, torch.Tensor]: + """The same sums in FP32: a true-FP32 matmul and a sum of squares. ``tile_n`` is ignored.""" + xf = x.float() + with true_fp32_matmul(): + mixes = xf @ fn.float().t() + return mixes, xf.square().sum(-1) + + def is_valid(self, x: torch.Tensor, fn: torch.Tensor, tile_n: int) -> None: + assert x.dim() == 2 and fn.dim() == 2 and fn.shape[1] == x.shape[1], ( + f"x is [M, K] and fn [mix, K]; got {tuple(x.shape)} and {tuple(fn.shape)}" + ) + assert x.is_contiguous() and fn.is_contiguous(), ( + "the kernel reads x and fn row-major with a row stride of K elements" + ) + + +mhc_gemm_sqrsum = MhcGemmSqrsum() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/mhc/mhc_post_mapping.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/mhc/mhc_post_mapping.py new file mode 100644 index 000000000000..e7e6e7f60922 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/mhc/mhc_post_mapping.py @@ -0,0 +1,125 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The mHC post-mapping: a sublayer's output mixed back into the residual streams. + +``torch.ops.trtllm.mhc_post_mapping`` (``mhcPostMappingKernel``: one CTA of 256 threads per row, eight +BF16 values per thread and step), the kernel ``mHC.post_mapping`` launches. GLM-5.3-Flash wraps every +sublayer in a hyper-connection of four residual streams. The target needs this op once per forward +step, after the last sublayer, where no next pre-mapping follows; ``mhc_fused_hc`` covers every other +boundary. + +For ``M`` rows of ``mult = 4`` streams of ``hidden`` values, with ``x`` the sublayer output, ``post`` +its post coefficients and ``comb`` its Sinkhorn-normalized stream combination:: + + out[t, j, :] = post[t, j] * x[t, :] + sum_k comb[t, k, j] * residual[t, k, :] + +that is ``post * x + comb^T @ residual``: row ``k`` of ``comb`` is a source stream and column ``j`` a +destination stream. The kernel accumulates in FP32 in that order (the ``post`` product, then +``k = 0..3`` as fused multiply-adds) and rounds each output to BF16 once. +""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The stream count the kernel is compiled for, and GLM-5.3-Flash's hidden size. +MULT = 4 +HIDDEN = 4096 + + +class MhcPostMapping(OpWrapper): + """``post * x + comb^T @ residual`` per row, rounded to BF16: the streams after a sublayer. + + Args: + x: ``[M, hidden]`` BF16, contiguous: the sublayer's output, after its TP reduction. + residual: ``[M, 4, hidden]`` BF16, contiguous: the streams the sublayer's pre-mapping read. + post: ``[M, 4]`` or ``[M, 4, 1]`` FP32, contiguous: the sublayer's post coefficients, + ``post_mult * sigmoid(...)`` from its pre-mapping (in ``[0, 2]`` for GLM-5.3-Flash). + comb: ``[M, 4, 4]`` FP32, contiguous: the sublayer's Sinkhorn-normalized combination, + indexed ``[source stream, destination stream]``. + + Returns: + ``[M, 4, hidden]`` BF16, a new contiguous tensor of ``residual``'s shape. + """ + + ARCHS = frozenset({Arch.SM_100}) + + _GLM = dict(mult=MULT, hidden=HIDDEN) + CELLS: tuple[Cell, ...] = ( + Cell(why="decode of one request without MTP: one CTA", spec=dict(rows=1, **_GLM)), + Cell(why="MTP-3 verify of one request (4 rows)", spec=dict(rows=4, **_GLM)), + Cell(why="MTP-3 verify of two requests (8 rows)", spec=dict(rows=8, **_GLM)), + Cell(why="MTP-3 verify of 8 requests (32 rows)", spec=dict(rows=32, **_GLM)), + Cell( + why="MTP-3 verify of 64 requests (256 rows): more CTAs than B200 has SMs", + spec=dict(rows=256, **_GLM), + ), + Cell( + why="a quarter of a 13,312-token prompt: one rank's rows under sequence parallelism", + spec=dict(rows=3328, **_GLM), + ), + ) + + note = """ + The kernel is launched without programmatic dependent launch and without a zero-row guard: a + launch of no rows leaves an invalid-configuration error pending for the next CUDA call, so + `is_valid` refuses it. The output is allocated per call, from the graph's private pool under CUDA + graph capture. `x` and `post` / `comb` are read as given: the wrapper neither casts nor copies + them. + """ + + def __call__( + self, + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + ) -> torch.Tensor: + rows, hidden = x.shape + out = torch.empty_like(residual, memory_format=torch.contiguous_format) + torch.ops.trtllm.mhc_post_mapping(residual, x, post, comb, out, rows, hidden) + return out + + def reference( + self, + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + ) -> torch.Tensor: + """FP32 in the kernel's order (the ``post`` product, then sources ``k = 0..3``), rounded once. + + Element-wise products and sums only, so no matmul precision setting reaches it. + """ + rows, hidden = x.shape + streams = residual.float().reshape(rows, MULT, hidden) + coeff = comb.float().reshape(rows, MULT, MULT) + out = post.float().reshape(rows, MULT, 1) * x.float().reshape(rows, 1, hidden) + for k in range(MULT): + out += coeff[:, k, :, None] * streams[:, k, None, :] + return out.to(torch.bfloat16).reshape(residual.shape) + + def is_valid( + self, + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + ) -> None: + assert x.dim() == 2, f"x is [M, hidden]; got {tuple(x.shape)}" + rows, hidden = x.shape + assert rows >= 1, "a launch of no rows leaves an error pending for the next CUDA call" + assert hidden % 8 == 0, "each thread moves 8 BF16 values per step; the kernel has no tail" + assert residual.numel() == rows * MULT * hidden and residual.shape[-1] == hidden, ( + f"residual must hold [{rows}, {MULT}, {hidden}]; got {tuple(residual.shape)}" + ) + assert post.numel() == rows * MULT and comb.numel() == rows * MULT * MULT, ( + f"post holds [{rows}, {MULT}] and comb [{rows}, {MULT}, {MULT}]" + ) + for name, t in (("x", x), ("residual", residual), ("post", post), ("comb", comb)): + assert t.is_contiguous(), f"{name} is read with dense row-major strides" + + +mhc_post_mapping = MhcPostMapping() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/moe/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/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/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/moe/decode_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/moe/decode_moe.py new file mode 100644 index 000000000000..4d7185f683ee --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/moe/decode_moe.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MoE branch of a decode step of at most 4 rows in one persistent CUDA launch (after the router GEMM). + +Every CTA (one per SM) computes the routing itself (sigmoid scores, FP32 correction bias, top-8, +normalized weights times the routed scaling factor) and derives the same static schedule from it, +while its producer warp already streams the shared expert, whose weights do not depend on the +routing. FC1 tiles (4 up and 4 gate rows of one expert over H, 32 KB) go round robin over the CTAs, +the CTAs without FC2 rows taking the tail; the compute warps multiply them with the rows, converted +to f16, on the tensor cores (FP32 accumulation, each 128 x 128 block's partial times its scale) and +the epilogue warps apply the clamped SwiGLU and the routing weight and publish the activation. e4m3 +converts to f16 exactly; a BF16 element does for magnitudes in [2^-14, 65504], while a smaller one +rounds to f16's subnormal grid (an absolute error of at most 2^-25) and a larger one overflows. +Both callers pass RMS-normalized layer inputs (the decoder layer's and the MTP layer's), far inside +that range. For FC2, each of 128 CTAs owns 32 output columns and streams every distinct expert's +down-projection tile, adding the expert's activation once its FC1 tiles are complete, so FC2 +overlaps the FC1 tail; the experts are summed in a fixed order in FP32 before the single BF16 +rounding. The rank's result is the moe_tp partial the caller all-reduces. The kernel is in +``csrc/decodeMoe.cu``. + +The BF16 activations are never quantized: the FC1 -> FC2 activations travel as an f16 hi / lo +pair, so this computes the reference product more closely than the FP8-activation grouped GEMM +it replaces at these sizes. With programmatic dependent launch (``TRTLLM_ENABLE_PDL``) only the +weights are read before the grid-dependency wait. +""" + +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 ... import native_ops + +MAX_ROWS = 4 +_TOPK = 8 + + +class DecodeMoe(OpWrapper): + """``scaling * sum_i w_i * expert_{e_i}(x) + shared(x)`` for each of at most 4 rows. + + ``logits`` are the router's FP32 ``[M, E]`` outputs; ``bias`` the FP32 correction bias. + ``w31`` / ``s31`` are the routed experts' ``[E, 2I, H]`` FP8 ``[up; gate]`` weights and + ``[E, 2I/128, H/128]`` scales, ``w2`` / ``s2`` the ``[E, H, I]`` down weights and scales; + ``sw31``, ``ss31``, ``sw2``, ``ss2`` the shared expert in the same per-expert layout. + Returns BF16 ``[M, H]``. + """ + + ARCHS = frozenset({Arch.SM_100}) + + # GLM-5.3-Flash at moe_tp4: hidden 4096, intermediate 2048 / 4 = 512, 288 experts. + CELLS: tuple[Cell, ...] = ( + Cell(why="MTP draft row", spec=dict(m=1, e=288, h=4096, i=512)), + Cell(why="MTP-3 verify of one request", spec=dict(m=4, e=288, h=4096, i=512)), + ) + + note = """ + Top-8 ties go to the lower expert index. The scores are computed as the routing kernels of the + larger steps' grouped GEMMs compute them (the tanh form of the sigmoid, plus the FP32 bias), so + near-equal experts order alike on both sides of the 4-row switch. The kernel is compiled for + H = 4096, I = 512 and E = 288 (moe_tp4) and holds at most 4 rows; larger steps are the caller's + to route to the grouped GEMMs. Its FC1 -> FC2 scratch and counters are device-global (the + counters double-buffered by call parity), so calls on one device must be ordered on one stream. + Its CTAs, one per SM, wait on each other, so they must all be resident at once: the op raises + when the device cannot hold the grid, and work that keeps SMs busy while waiting on this launch + would hang it. Each distinct expert of the step is read from HBM once, so the time follows the + number of distinct experts the rows select (one request's MTP-3 verify rows choose 14-33 of + them, shared expert included). + """ + + def __call__(self, x, logits, bias, scaling, limit, w31, s31, w2, s2, sw31, ss31, sw2, ss2): + out = torch.empty(x.shape, dtype=torch.bfloat16, device=x.device) + native_ops.decode_moe( + x, + logits, + bias, + scaling, + limit, + w31, + s31, + w2, + s2, + sw31, + ss31, + sw2, + ss2, + out, + get_env_enable_pdl(), + ) + return out + + def reference(self, x, logits, bias, scaling, limit, w31, s31, w2, s2, sw31, ss31, sw2, ss2): + def dequant(w, s): + r, c = w.shape + return (w.float().view(r // 128, 128, c // 128, 128) * s[:, None, :, None]).view(r, c) + + def expert(w_up_gate, s_up_gate, w_down, s_down, row): + h = dequant(w_up_gate, s_up_gate) @ row + inter = h.shape[0] // 2 + u, g = h[:inter].clamp(-limit, limit), h[inter:].clamp(max=limit) + return dequant(w_down, s_down) @ (torch.nn.functional.silu(g) * u) + + out = [] + with true_fp32_matmul(): + score = torch.sigmoid(logits.float()) + ids = torch.topk(score + bias.float(), _TOPK, dim=-1).indices + wts = score.gather(1, ids) + wts = wts / wts.sum(-1, keepdim=True) * scaling + for t in range(x.shape[0]): + row = x[t].float() + acc = expert(sw31, ss31, sw2, ss2, row) + for i in range(_TOPK): + e = int(ids[t, i]) + acc = acc + wts[t, i] * expert(w31[e], s31[e], w2[e], s2[e], row) + out.append(acc) + return torch.stack(out).to(torch.bfloat16) + + def is_valid(self, x, logits, bias, scaling, limit, w31, s31, w2, s2, sw31, ss31, sw2, ss2): + m, h = x.shape + assert 1 <= m <= MAX_ROWS, f"{m} rows; at most {MAX_ROWS}" + assert x.dtype == torch.bfloat16 and x.is_contiguous() + assert ( + logits.dtype == torch.float32 and logits.is_contiguous() and bias.dtype == torch.float32 + ) + for t in (w31, w2, sw31, sw2): + assert t.dtype == torch.float8_e4m3fn and t.is_contiguous() + for t in (s31, s2, ss31, ss2): + assert t.dtype == torch.float32 and t.is_contiguous() + assert h == 4096 and logits.shape[1] == 288 and w2.shape[2] == 512, ( + "the kernel is compiled for H 4096, E 288, I 512" + ) + + +decode_moe = DecodeMoe() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/moe/fp8_block_scale_moe_runner.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/moe/fp8_block_scale_moe_runner.py new file mode 100644 index 000000000000..cf792b287c79 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/moe/fp8_block_scale_moe_runner.py @@ -0,0 +1,436 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DeepSeek-style FP8 block-scale mixture-of-experts layer, routed in the kernel: sigmoid + +correction-bias top-k routing, grouped FC1 GEMM, FP8 requantization, clamped SwiGLU, FP8 +requantization, grouped FC2 GEMM and the routing-weighted combine, in one trtllm-gen call. + +Copied from qwen4_exp's entry (``models/qwen4_exp/ops/moe/fp8_block_scale_moe_runner.py``), which +wraps the same op. The reference here diverges from that one: it also routes from the logits +(``route_deepseek_v3``) where qwen4_exp's models the pre-routed path only; it rounds each expert's +FC2 output and the combine weights to BF16, as the kernel reads them, where qwen4_exp's combines in +FP32; and ``requantize_1x128`` divides tensor by tensor rather than by a Python scalar (a rounded +reciprocal times the value). The cells are at GLM-5.3-Flash's TP4 geometry. +""" + +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, + true_fp32_matmul, +) + +#: Block edge on every scale axis: 1x128 activations, 128x128 weights. +BLOCK = 128 + +#: Largest finite e4m3 magnitude. +E4M3_MAX = 448.0 + +#: The activation kernel's floor on a block amax (`kDsActAmaxEpsilon` in +#: trtllmGen/DevKernel.cu). Only an all-zero block reaches it. +ACT_AMAX_FLOOR = 1.0e-10 + +#: `routing_method_type` values (`RoutingMethodType` in kernels/moe/trtllmGen/runner.h). +ROUTING_RENORMALIZE = 1 +ROUTING_DEEPSEEK_V3 = 2 +ROUTING_MINIMAX2 = 5 + +#: The routing methods that read `routing_bias`; every other one ignores it. +_BIAS_READING_ROUTING = frozenset({ROUTING_DEEPSEEK_V3, ROUTING_MINIMAX2}) + + +def dequantize_block_weight(codes: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + """One expert's `[R, C]` fp32 weight from e4m3 codes and `[R/128, C/128]` scales.""" + r, c = codes.shape + blocks = codes.float().reshape(r // BLOCK, BLOCK, c // BLOCK, BLOCK) + return (blocks * scales.float().reshape(r // BLOCK, 1, c // BLOCK, 1)).reshape(r, c) + + +def requantize_1x128(x: torch.Tensor) -> torch.Tensor: + """The e4m3 round trip the kernel applies to the FC1 output and to the activation. + + Per row and per 128 consecutive columns: `scale = max(amax, 1e-10) / 448`, then + round-to-nearest-even into e4m3 saturating at +-448. The plain `amax / 448` scale, not a + power of two. + """ + shape = x.shape + blocks = x.reshape(*shape[:-1], shape[-1] // BLOCK, BLOCK) + amax = blocks.abs().amax(dim=-1, keepdim=True).clamp(min=ACT_AMAX_FLOOR) + # Tensor-by-tensor divisions: torch divides by a Python scalar as a multiply by its + # rounded reciprocal. + scale = amax / torch.full_like(amax, E4M3_MAX) + codes = (blocks / scale).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn) + return (codes.float() * scale).reshape(shape) + + +def route_deepseek_v3( + routing_logits: torch.Tensor, + routing_bias: torch.Tensor, + top_k: int, + routed_scaling_factor: Optional[float], +) -> tuple[torch.Tensor, torch.Tensor]: + """Ungrouped DeepSeek-V3 routing as the kernel evaluates it: `(ids [T, k] int64, weights + [T, k] bf16)`. + + `s = sigmoid(logit)` in its tanh form, `0.5 * tanh(x / 2) + 0.5`, in fp32. The experts are + chosen by `s + bias`; each chosen expert's weight is its unbiased score, recovered as + `(s + bias) - bias` the way the kernel recovers it, times `routed_scaling_factor` (1 when + None) over the sum of the chosen scores, rounded to bf16 -- the dtype the combine reads. + Ties go to the lower expert index here; the kernel's tie rule is not certified. + """ + score = 0.5 * torch.tanh(0.5 * routing_logits.float()) + 0.5 + bias = routing_bias.float() + choice = score + bias + ids = torch.sort(choice, dim=-1, descending=True, stable=True).indices[:, :top_k] + picked = choice.gather(1, ids) - bias[ids] + scale = 1.0 if routed_scaling_factor is None else routed_scaling_factor + weights = picked * scale / picked.sum(dim=-1, keepdim=True) + return ids, weights.to(torch.bfloat16) + + +class Fp8BlockScaleMoeRunner(OpWrapper): + """One FP8 block-scale MoE layer over e4m3 activations, routed from the router logits. + + Routing, per token `t` (when `topk_ids` is None, `routing_method_type` DeepSeek-V3 and + `n_group` at most 1): `s = sigmoid(routing_logits[t])`, the `top_k` experts by `s + + routing_bias`, and combine weights `s_e * routed_scaling_factor / sum(s_chosen)` rounded to + bf16 -- see `route_deepseek_v3`. With `topk_ids` / `topk_weights` given instead, those are + the routing, used as given. + + Then per token `t` and slot `j`, with `A[t]` the fp32 activations + (`hidden_states[t, k] * hidden_states_scale[k // 128, t]`) and `W_*` the dequantized + 128x128-block weights of local expert `e = id[t, j] - local_expert_offset` (slots outside + `[0, local_num_experts)` contribute nothing): + + up, gate = requant(A[t] @ W1[e].T).split(I) # FC1 rows are [up ; gate] + gate, up = min(gate, L), clamp(up, -L, L) # L = gemm1_clamp_limit + y = bf16(requant(silu(gate) * up) @ W2[e].T) + out[t] += w[t, j] * y + + accumulated in fp32 and stored bf16, `requant` being `requantize_1x128`. FC1 rows are + `[up ; gate]` -- up first, the checkpoint's `[w3 ; w1]` -- the opposite of a dense MLP's + fused `[gate ; up]` weight. + + Fusion boundary: routing, permutation, both GEMMs with on-the-fly dequantization, both + intermediate requantizations, the clamped SwiGLU and the weighted combine are inside. + The router GEMM, quantizing the bf16 hidden states (`fp8_quantize_1x128`), the shared + expert, any all-reduce and the residual are the caller's. + + `hidden_states_scale` is `[H/128, T]` -- K-block-major, exactly what `fp8_quantize_1x128` + returns on this architecture. Pass it straight through: at `T == H/128` (32 rows at + H 4096) its transpose has the same shape and passes every check, and the answer is wrong. + + Returns a fresh `[T, H]` bf16 tensor, or an empty `[0]` bf16 tensor when `output` is given + (the result is then in `output`). + """ + + ARCHS = frozenset({Arch.SM_100}) + + #: GLM-5.3-Flash's routed experts at TP4: hidden 4096, moe_intermediate_size 2048 split + #: four ways (512), all 288 experts local, ungrouped DeepSeek-V3 top-8 routing scaled by + #: 2.5, swiglu_limit 10. + GLM53_FLASH_TP4 = dict( + hidden=4096, + inter=512, + num_experts=288, + top_k=8, + n_group=1, + topk_group=1, + routed_scaling_factor=2.5, + gemm1_clamp_limit=10.0, + local_expert_offset=0, + local_num_experts=288, + ) + + # Only GLM-5.3-Flash calls this entry. Its own cells are the routing-kernel boundaries at + # GLM's geometry -- 288 experts dispatch to the 384-expert tier, where 2 to 8 rows run the + # small-token kernel, 9 to 256 the single-cluster kernel and more the multi-kernel path. + # The row counts the target runs are in the test. + CELLS: tuple[Cell, ...] = ( + Cell( + why="9 rows: the first count on the single-cluster routing kernel", + spec=dict(tokens=9, **GLM53_FLASH_TP4), + ), + Cell( + why="2,047 rows: past the single-cluster kernel's 256, the multi-kernel routing path", + spec=dict(tokens=2047, **GLM53_FLASH_TP4), + ), + ) + + note = """ + Metadata consumed: none. Two process-global caches sit behind the call: the autotuner's + profiling cache and the C++ runner cache. Outside an `autotune()` context a cold cache + runs a deterministic fallback tactic keyed by the row count; a warm, tuned tactic is not + certified here. + + The op is a Python `torch.library.custom_op` (`trtllm::fp8_block_scale_moe_runner` in + `_torch/custom_ops/trtllm_gen_custom_ops.py`) over the compiled-in C++ class + `torch.classes.trtllm.FP8BlockScaleMoERunner` (thop/moe/trtllmGen/fp8BlockScaleMoe.cpp), + which launches the routing kernels, trtllm-gen's precompiled batched-GEMM cubins, a + separate activation kernel and a finalize kernel: nothing is JIT-compiled. + + Routing. DeepSeek-V3 with `n_group` at most 1 runs the custom routing kernels with a + sigmoid + bias preprocess and a scaled sum-normalize postprocess (kernels/moe/trtllmGen/ + runner.cu); which kernel depends on the row count at the expert tier (288 experts use the + 384 tier): 1 row the cooperative block kernel, 2-8 the small-token kernel, 9-256 the + single-cluster kernel, more the histogram + offsets path. The sigmoid is the tanh form and + the weights are written as bf16, which `route_deepseek_v3` reproduces. Both + `routing_logits` and `routing_bias` may be fp32 or bf16; GLM-5.3-Flash passes fp32 for + both (rounding the correction bias to bf16 can reorder near-equal experts). The kernel's + tie rule is not certified: the test's routing has no near-ties. + + Preconditions, each enforced loudly by the op: `hidden_states` `[T, H]` e4m3 and + `hidden_states_scale` `[H/128, T]` fp32 (see the class docstring); weights in the plain + per-expert layout, no shuffle and no swizzle: `gemm1_weights` `[E, 2I, H]` e4m3 rows + `[up ; gate]`, `gemm1_weights_scale` `[E, 2I/128, H/128]` fp32, `gemm2_weights` `[E, H, I]`, + `gemm2_weights_scale` `[E, H/128, I/128]`, `E == local_num_experts`; H and I multiples of + 128; `intermediate_size == gemm1_weights.shape[1] // 2`; `routing_logits` `[T, num_experts]` + and `routing_bias` `[num_experts]`; `num_experts > top_k`. + + `gemm1_clamp_limit` is one float for every expert, applied after the FC1 requantization + (gate on the upper side, up on both). Any value given is applied, 0 included -- unlike + `silu_and_mul`, which drops a non-positive limit. None means no clamp. + + Silent on the raw op, refused by `is_valid`: any non-contiguous tensor (read as dense); + `act_type != 0` (the Python runner forces SwiGLU); `routing_logits` or `routing_bias` + alongside `topk_ids` (both ignored); a `routing_bias` with a routing method that does not + read it; DeepSeek-V3 routing from logits without a `routing_bias` (the routing kernels read + it for every expert with no null check -- read from the source, not driven); a + `local_expert_offset` outside `[0, num_experts)`. + + The registered fake returns `[T, 2H]` (it assumes packed fp4 hidden states): an upstream + defect; eager is unaffected. + + Certified (sm_100): GLM-5.3-Flash's routed experts at TP4 as listed above, routed in the + kernel, at 5, 8, 9, 32, 256 and 2,047 rows. + + Not certified: the pre-routed path (`topk_ids` given), grouped routing (`n_group > 1`), + other routing methods, bf16 logits or bias, `num_fused_shared_experts > 0`, + expert-parallel windows, `output=`, a warm autotuner cache, CUDA-graph capture, other + geometries, sm_103. + """ + + def __call__( + self, + 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, + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + num_experts: int, + top_k: int, + num_fused_shared_experts: Optional[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, + topk_weights: Optional[torch.Tensor] = None, + topk_ids: Optional[torch.Tensor] = None, + act_type: int = 0, + gemm1_clamp_limit: Optional[float] = None, + output: Optional[torch.Tensor] = None, + tune_max_num_tokens: int = 8192, + use_dp: bool = False, + ) -> torch.Tensor: + return torch.ops.trtllm.fp8_block_scale_moe_runner( + routing_logits, + routing_bias, + hidden_states, + hidden_states_scale, + gemm1_weights, + gemm1_weights_scale, + gemm2_weights, + gemm2_weights_scale, + num_experts, + top_k, + num_fused_shared_experts, + n_group, + topk_group, + intermediate_size, + local_expert_offset, + local_num_experts, + routed_scaling_factor, + routing_method_type, + topk_weights=topk_weights, + topk_ids=topk_ids, + act_type=act_type, + gemm1_clamp_limit=gemm1_clamp_limit, + output=output, + tune_max_num_tokens=tune_max_num_tokens, + use_dp=use_dp, + ) + + def reference( + self, + 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, + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + num_experts: int, + top_k: int, + num_fused_shared_experts: Optional[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, + topk_weights: Optional[torch.Tensor] = None, + topk_ids: Optional[torch.Tensor] = None, + act_type: int = 0, + gemm1_clamp_limit: Optional[float] = None, + output: Optional[torch.Tensor] = None, + tune_max_num_tokens: int = 8192, + use_dp: bool = False, + ) -> torch.Tensor: + """The layer over bit-identical operands: fp32 GEMMs, the kernel's two e4m3 + requantizations, its bf16 FC2 rows and bf16 combine weights, one bf16 store. + + Routes with `route_deepseek_v3` unless `topk_ids` / `topk_weights` are given. Not + modeled, so refused: fused shared experts, grouped routing and any routing method but + DeepSeek-V3 when routing from logits. Ignored: `num_experts` (the routing space is + `routing_logits.shape[1]`), `topk_group`, `act_type`, `output`, + `tune_max_num_tokens`, `use_dp`. Returns the `[T, H]` result even when `output` is + given. + """ + assert not num_fused_shared_experts, "the reference does not model fused shared experts" + if topk_ids is None: + assert routing_logits is not None and routing_bias is not None + assert routing_method_type == ROUTING_DEEPSEEK_V3 and (n_group or 1) <= 1, ( + "the reference routes ungrouped DeepSeek-V3 only" + ) + ids, weights = route_deepseek_v3( + routing_logits, routing_bias, top_k, routed_scaling_factor + ) + else: + assert topk_weights is not None + ids, weights = topk_ids.long(), topk_weights.to(torch.bfloat16) + weights = weights.float() + + num_tokens, hidden = hidden_states.shape + nb = hidden // BLOCK + act_in = ( + hidden_states.float().reshape(num_tokens, nb, BLOCK) + * hidden_states_scale.float().t().reshape(num_tokens, nb, 1) + ).reshape(num_tokens, hidden) + out = torch.zeros(num_tokens, hidden, dtype=torch.float32, device=hidden_states.device) + local = ids - local_expert_offset + live = local[(local >= 0) & (local < local_num_experts)] + inter = intermediate_size + with true_fp32_matmul(): + for e in torch.unique(live).tolist(): + tok, slot = (local == e).nonzero(as_tuple=True) + w1 = dequantize_block_weight(gemm1_weights[e], gemm1_weights_scale[e]) + h1 = requantize_1x128(act_in[tok] @ w1.t()) + up, gate = h1[:, :inter], h1[:, inter:] + if gemm1_clamp_limit is not None: + gate = gate.clamp(max=gemm1_clamp_limit) + up = up.clamp(-gemm1_clamp_limit, gemm1_clamp_limit) + act = requantize_1x128(gate / (1.0 + torch.exp(-gate)) * up) + w2 = dequantize_block_weight(gemm2_weights[e], gemm2_weights_scale[e]) + y = (act @ w2.t()).to(torch.bfloat16).float() + # A row picks an expert at most once, so `tok` holds no duplicates. + out[tok] += y * weights[tok, slot].unsqueeze(1) + return out.to(torch.bfloat16) + + def is_valid( + self, + 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, + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + num_experts: int, + top_k: int, + num_fused_shared_experts: Optional[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, + topk_weights: Optional[torch.Tensor] = None, + topk_ids: Optional[torch.Tensor] = None, + act_type: int = 0, + gemm1_clamp_limit: Optional[float] = None, + output: Optional[torch.Tensor] = None, + tune_max_num_tokens: int = 8192, + use_dp: bool = False, + ) -> None: + """Inputs the op accepts and answers wrongly, or faults on. + + The transposed activation-scale buffer at `T == H/128` is one more, and no metadata + check can see it. + """ + # The kernels take raw data pointers: a strided view is read as 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), + ("topk_weights", topk_weights), + ("topk_ids", topk_ids), + ("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 Python runner overwrites the activation with SwiGLU. + assert act_type == 0, ( + f"act_type={act_type} is silently ignored -- this op always computes SwiGLU" + ) + if topk_ids is not None: + # Once topk_ids carry the routing, nothing reads the logits or the bias. + assert routing_logits is None and routing_bias is None, ( + "routing_logits and routing_bias are ignored when topk_ids/topk_weights are " + "given (routing has already happened)" + ) + else: + # DeepSeek-V3 routing loads routing_bias[e] for every expert with no null check. + if routing_method_type == ROUTING_DEEPSEEK_V3: + assert routing_bias is not None, ( + "DeepSeek-V3 routing from logits reads routing_bias for every expert and " + "faults on None; pass zeros for no correction" + ) + # Only the DeepSeek-V3 and MiniMax-2 routing kernels are handed the bias. + if routing_bias is not None: + assert routing_method_type in _BIAS_READING_ROUTING, ( + f"routing_method_type={routing_method_type} ignores routing_bias; only " + "DeepSeek-V3 (2) and MiniMax-2 (5) routing read it" + ) + # `num_experts` is the routing space, not the local window. + assert 0 <= local_expert_offset < num_experts, ( + f"local_expert_offset={local_expert_offset} is outside the routing space " + f"[0, num_experts={num_experts}); num_experts is every routed expert, not the " + "window's count" + ) + + +fp8_block_scale_moe_runner = Fp8BlockScaleMoeRunner() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/moe/grouped_blockwise_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/moe/grouped_blockwise_gemm.py new file mode 100644 index 000000000000..38b16b29621c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/moe/grouped_blockwise_gemm.py @@ -0,0 +1,2143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# The kernel body is adapted from tensorrt_llm/_torch/kernels/blackwell/blockwise_gemm/ +# blockwise_gemm.py, itself derived from the CUTLASS CuTe DSL blockwise GEMM example: +# +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +"""Grouped (contiguous MoE layout) FP8 GEMM with FP32 block scales on SM100, in CuTe DSL. + +This is TensorRT-LLM's persistent, warp-specialized dense blockwise GEMM +(``Sm100BlockwiseGemmKernel``) with the batch mode turned into the expert index. The FP32 +scales are applied by promotion, exactly as in the dense kernel, so arbitrary (not just +power-of-two) 1 x 128 activation and 128 x 128 weight scales are kept: the result is as +precise as the dense kernel's. + +Layouts (``rows`` is the worst-case row capacity, a multiple of 128): + +* ``a`` ``[rows, K]`` FP8 E4M3, row-major. Rows are grouped by expert in contiguous segments, + each starting at a multiple of 128. +* ``sfa`` ``[K / 128, rows]`` FP32: K-block-major, the layout ``fp8_quantize_1x128`` returns. +* ``b`` ``[E, N, K]`` FP8 E4M3, K-major per expert; ``sfb`` ``[E, N / 128, K / 128]`` FP32. +* ``tile_expert`` ``[rows / 128]`` Int32 on the GPU: the expert of each 128-row tile, -1 if + unused. +* ``out`` ``[rows, N]`` BF16. + +The tile scheduler walks every (M tile, N tile) of the worst case, so the grid never depends +on device data. The scheduler warp, which already broadcasts each tile's coordinates to the +other warps through shared memory, reads ``tile_expert`` for the tile, skips tiles whose +expert is -1 (no loads, MMA or store) and publishes the expert in the batch slot. The TMA and +scale warps then index B and SFB by that expert, and A, SFA and C by batch 0. + +A cluster must work on one expert, so the MMA tile is at most 128 rows and the cluster grows +only along N (a 2-CTA MMA uses 64-row CTA tiles inside one 128-row group). + +``grouped_blockwise_fp8_gemm_gathered`` reads A as ``x[row_index]`` from the un-permuted +activations instead of a pre-gathered copy. Each lane of the TMA warp issues one +``tile::gather4`` (four source rows of 128 bytes) per K tile into the SW128 A stage, through a +2D descriptor whose box is one row; with an A-multicast cluster the CTAs split the row groups. +The scale warp copies each row's scales with 4-byte ``cp.async`` from the same indices. The +bytes in shared memory match the TMA tile load's, so the output is bitwise the same. +""" + +from typing import Callable, Dict, Tuple, Union + +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 +import torch +from cutlass._mlir.dialects import cute as _cute_ir +from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.runtime import make_ptr +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +from tensorrt_llm._torch.kernels.blackwell.blockwise_gemm.blockwise_gemm import ( + Sm100BlockwiseGemmKernel, +) + + +@dsl_user_op +def _tma_gather4_load(tma_atom, smem_dst, mbar, col, rows, mcast_mask, *, loc=None, ip=None): + """Issue one 1-CTA ``cp.async.bulk.tensor.2d...tile::gather4``: rows ``rows[0..3]``, + columns ``[col, col + box)``, into four consecutive smem box rows at ``smem_dst``. + + A row index outside the tensor (e.g. -1) makes the TMA unit fill that row with zeros. + With ``mcast_mask`` the rows land at the same offset in every CTA of the mask. + """ + exec_atom = _cute_nvgpu_ir.atom_make_exec_tma(tma_atom._trait.value, loc=loc, ip=ip) + desc_ptr_ty = _cute_ir.PtrType.get( + _cute_nvgpu_ir.TmaDescriptorTiledType.get(), cute.AddressSpace.generic, 64 + ) + desc = _cute_nvgpu_ir.get_tma_desc_addr(desc_ptr_ty, exec_atom, loc=loc, ip=ip) + operands = [ + cutlass.Int32(smem_dst.toint()).ir_value(loc=loc, ip=ip), + desc.toint().ir_value(loc=loc, ip=ip), + cutlass.Int32(mbar.toint()).ir_value(loc=loc, ip=ip), + cutlass.Int32(col).ir_value(loc=loc, ip=ip), + ] + [cutlass.Int32(rows[j]).ir_value(loc=loc, ip=ip) for j in range(4)] + if mcast_mask is None: + asm = ( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes [$0], [$1, {$3, $4, $5, $6, $7}], [$2];" + ) + constraints = "r,l,r,r,r,r,r,r" + else: + asm = ( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [$0], [$1, {$3, $4, $5, $6, $7}], [$2], $8;" + ) + operands.append(cutlass.Int16(mcast_mask).ir_value(loc=loc, ip=ip)) + constraints = "r,l,r,r,r,r,r,r,h" + llvm.inline_asm( + None, + operands, + asm, + constraints, + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +class Sm100GroupedBlockwiseGemmKernel(Sm100BlockwiseGemmKernel): + """``Sm100BlockwiseGemmKernel`` over a contiguous grouped layout, B indexed per M tile. + + :param use_2cta_instrs: Use the 2-CTA MMA (then the MMA tile M must be 128). + :param mma_tiler_mn: MMA tile; M in {64, 128}, N 128. + :param cluster_shape_mn: (1, cn), or (2, cn) with the 2-CTA MMA; cn a power of two. + :param raster_along_m: Walk tiles M first. N first (the default) keeps the concurrently + running CTAs on the same A rows, which lets A come from L2. + """ + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + sfa: cute.Tensor, + sfb: cute.Tensor, + tile_expert: cute.Tensor, + row_index: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + """Set up the TMA atoms, the grid and the shared storage, then launch the kernel. + + ``a`` is (rows, K, 1), ``b`` is (N, K, E), ``c`` is (rows, N, 1), ``sfa`` is + (rows, K / 128, 1), ``sfb`` is (N / 128, K / 128, E) and ``tile_expert`` is + (rows / 128,) Int32. With ``gather_a``, ``a`` and ``sfa`` are the un-permuted + (T, K, 1) and (T, K / 128, 1) tensors and ``row_index`` (rows,) Int32 gives each + output row's source row; otherwise ``row_index`` is unused. + """ + self.a_dtype = a.element_type + self.b_dtype = b.element_type + self.c_dtype = c.element_type + self.sfa_dtype = sfa.element_type + self.sfb_dtype = sfb.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + + # Check if input data types are compatible with MMA instruction + if cutlass.const_expr(self.a_dtype != self.b_dtype): + raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.a_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + if cutlass.const_expr(self.gather_a): + # A 2D descriptor whose box is one 128-byte row with the 128-byte swizzle of the + # K-major A tile: a tile::gather4 through it writes four rows at the offsets the + # MMA reads (row m of a stage at m * 128 bytes). + a_row_layout = cute.make_composed_layout( + cute.make_swizzle(3, 4, 3), + 0, + cute.make_layout((1, self.mma_tiler[2]), stride=(self.mma_tiler[2], 1)), + ) + tma_atom_a, tma_tensor_a = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + a[(None, None, 0)], + a_row_layout, + (1, self.mma_tiler[2]), + ) + else: + a_op = self._get_tma_atom_kind(atom_thr_size, self.is_a_mcast) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=(cutlass.TFloat32 if a.element_type is cutlass.Float32 else None), + ) + + # Setup TMA load for B + b_op = self._get_tma_atom_kind(atom_thr_size, self.is_b_mcast) + 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, + internal_type=(cutlass.TFloat32 if b.element_type is cutlass.Float32 else None), + ) + + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + self.num_tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size + + # Setup TMA store for C + tma_atom_c = None + tma_tensor_c = None + c_cta_v_layout = cute.composition(cute.make_identity_layout(c.shape), self.epi_tile) + epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + c, + epi_smem_layout, + c_cta_v_layout, + ) + + tensor_sfa = cute.make_tensor( + sfa.iterator, + cute.make_layout( + ( + (self.scale_granularity_m, sfa.shape[0]), + (self.scale_granularity_k, sfa.shape[1]), + sfa.shape[2], + ), + stride=( + (0, sfa.layout.stride[0]), + (0, sfa.layout.stride[1]), + sfa.layout.stride[2], + ), + ), + ) + tensor_sfb = cute.make_tensor( + sfb.iterator, + cute.make_layout( + ( + (self.scale_granularity_n, sfb.shape[0]), + (self.scale_granularity_k, sfb.shape[1]), + sfb.shape[2], + ), + stride=( + (0, sfb.layout.stride[0]), + (0, sfb.layout.stride[1]), + sfb.layout.stride[2], + ), + ), + ) + + # Compute grid size + self.tile_sched_params, grid = self._compute_grid_grouped( + c, self.cta_tile_shape_mnk, self.cluster_shape_mn, max_active_clusters + ) + + self.buffer_align_bytes = 1024 + + c_smem_size = cute.cosize(self.c_smem_layout_staged.outer) + + # Define shared storage for kernel + @cute.struct + class SharedStorage: + # (bidx, bidy, bidz, valid) + sInfo: cute.struct.Align[ + cute.struct.MemRange[cutlass.Int32, 4 * self.num_tile_stage], + # 1 byte alignment + 1, + ] + ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + scale_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_scale_stage * 2] + acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] + tile_info_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_tile_stage * 2] + epi_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 1 * 2] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC: cute.struct.Align[ + cute.struct.MemRange[ + self.c_dtype, + c_smem_size, + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_M, MMA_K, STAGE) + sA: cute.struct.Align[ + cute.struct.MemRange[self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + # (MMA, MMA_N, MMA_K, STAGE) + sB: cute.struct.Align[ + cute.struct.MemRange[self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + # (granularity_m, repeat_m), (granularity_k, repeat_k), num_scale_stage) + sSFA: cute.struct.Align[ + cute.struct.MemRange[self.sfa_dtype, cute.cosize(self.sfa_smem_layout_staged)], + self.buffer_align_bytes, + ] + # (granularity_n, repeat_n), (granularity_k, repeat_k), num_scale_stage) + sSFB: cute.struct.Align[ + cute.struct.MemRange[self.sfb_dtype, cute.cosize(self.sfb_smem_layout_staged)], + self.buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + # Launch the kernel synchronously + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c, + tensor_sfa, + tensor_sfb, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.c_smem_layout_staged, + self.sfa_smem_layout_staged, + self.sfb_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + tile_expert, + row_index, + sfa, + epilogue_op, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + smem=self.shared_storage.size_in_bytes(), + stream=stream, + min_blocks_per_mp=1, + ) + return + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + mSFA_mkl: cute.Tensor, + mSFB_nkl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + mTileExpert: cute.Tensor, + mRowIndex: cute.Tensor, + mSFX: cute.Tensor, + epilogue_op: cutlass.Constexpr, + ): + """Persistent grouped GEMM. The L coordinate a warp reads from the tile info is the + tile's expert: it indexes B and SFB, while A, C and SFA always use L = 0.""" + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + lane_idx = cute.arch.lane_idx() + + # + # Prefetch tma desc + # + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + cpasync.prefetch_descriptor(tma_atom_c) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + # + # Setup cta/thread coordinates + # + # Coords inside cluster + 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) + # Coord inside cta + tidx, _, _ = cute.arch.thread_idx() + + # + # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier + # + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + tmem_dealloc_mbar_ptr = storage.tmem_dealloc_mbar_ptr + tmem_holding_buf = storage.tmem_holding_buf + + # Initialize mainloop ab_pipeline (barrier) and states + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + ab_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Initialize mainloop scale_pipeline (barrier) and states + scale_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * 1, + ) + scale_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.epilog_warp_id), + ) + scale_pipeline = pipeline.PipelineCpAsync.create( + barrier_storage=storage.scale_mbar_ptr.data_ptr(), + num_stages=self.num_scale_stage, + producer_group=scale_pipeline_producer_group, + consumer_group=scale_pipeline_consumer_group, + defer_sync=True, + ) + + # Initialize acc_pipeline (barrier) and states + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilog_warp_id) * (2 if use_2cta_instrs else 1) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Initialize epilogue pipeline (barrier) and states + epi_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.acc_update_warp_id), + ) + epi_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.epilog_warp_id), + ) + epi_pipeline = pipeline.PipelineAsync.create( + barrier_storage=storage.epi_mbar_ptr.data_ptr(), + num_stages=1, + producer_group=epi_pipeline_producer_group, + consumer_group=epi_pipeline_consumer_group, + defer_sync=True, + ) + + # Initialize tile info pipeline (barrier) and states + tile_info_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * 1, + ) + tile_info_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_wo_sched, + ) + tile_info_pipeline = pipeline.PipelineAsync.create( + barrier_storage=storage.tile_info_mbar_ptr.data_ptr(), + num_stages=self.num_tile_stage, + producer_group=tile_info_pipeline_producer_group, + consumer_group=tile_info_pipeline_consumer_group, + defer_sync=True, + ) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + tmem_holding_buf, + barrier_for_retrieve=self.tmem_alloc_barrier, + allocator_warp_id=self.epilog_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=tmem_dealloc_mbar_ptr, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # + # Setup smem tensor A/B/C/Scale + # + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC = storage.sC.get_tensor(c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner) + # (MMA, MMA_M, MMA_K, STAGE) + sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner) + # (MMA, MMA_N, MMA_K, STAGE) + sB = storage.sB.get_tensor(b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner) + # (granularity_m, repeat_m), (granularity_k, repeat_k), num_scale_stage) + sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged) + # (granularity_n, repeat_n), (granularity_k, repeat_k), num_scale_stage) + sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged) + # (bidx, bidy, bidz, valid) + info_layout = cute.make_layout((4, self.num_tile_stage), stride=(1, 4)) + sInfo = storage.sInfo.get_tensor(info_layout) + + # + # Compute multicast mask for A/B buffer full + # + a_full_mcast_mask = None + b_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + # + # Local_tile partition global tensors + # + # (bN, bK, loopN, loopK, loopL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + # (bM, bN, loopM, loopN, loopL) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + # (bM, bK, loopM, loopK, loopL) + gSFA_mkl = cute.local_tile( + mSFA_mkl, + cute.slice_(self.cta_tile_shape_mnk, (None, 0, None)), + (None, None, None), + ) + # (bN, bK, loopN, loopK, loopL) + gSFB_nkl = cute.local_tile( + mSFB_nkl, + cute.slice_(self.cta_tile_shape_mnk, (0, None, None)), + (None, None, None), + ) + # coordinate + cSFA_mkl = cute.make_identity_tensor(cute.shape(mSFA_mkl)) + cSFB_nkl = cute.make_identity_tensor(cute.shape(mSFB_nkl)) + # (bM, bK, loopM, loopK, loopL) + cSFA = cute.local_tile( + cSFA_mkl, + cute.slice_(self.cta_tile_shape_mnk, (None, 0, None)), + (None, None, None), + ) + # (bN, bK, loopN, loopK, loopL) + cSFB = cute.local_tile( + cSFB_nkl, + cute.slice_(self.cta_tile_shape_mnk, (0, None, None)), + (None, None, None), + ) + k_tile_cnt = cute.size(gB_nkl, mode=[3]) + + # + # Partition global tensor for TiledMMA_A/B/C + # + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + # (MMA, MMA_N, MMA_K, loopN, loopK, loopL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_M, MMA_N, loopM, loopN, loopL) + tCgC = thr_mma.partition_C(gC_mnl) + + # scale viewed as C tensor + sSFA_view_as_C_layout = cute.make_layout( + ( + (self.scale_granularity_m, self.scale_m_per_tile), + self.cta_tile_shape_mnk[1], + self.num_scale_stage, + ), + stride=((0, 1), 0, self.scale_m_per_tile), + ) + sSFB_view_as_C_layout = cute.make_layout( + ( + self.cta_tile_shape_mnk[0], + (self.scale_granularity_n, self.scale_n_per_tile), + self.num_scale_stage, + ), + stride=(0, (0, 1), self.scale_n_per_tile), + ) + sSFA_view_as_C = cute.make_tensor(sSFA.iterator, sSFA_view_as_C_layout) + sSFB_view_as_C = cute.make_tensor(sSFB.iterator, sSFB_view_as_C_layout) + + # + # Partition global/shared tensor for TMA load A/B + # + # TMA load A partition_S/D + a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tAsA = None + tAgA = None + if cutlass.const_expr(not self.gather_a): + # (bM, bK, loopM, loopK, loopL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (MMA, MMA_M, MMA_K, loopM, loopK, loopL) + tCgA = thr_mma.partition_A(gA_mkl) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + # TMA load B partition_S/D + b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # + # Partition global/shared tensor for TMA load A/B + # + # load scaleA/scaleB + atom_copy = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(), + mSFA_mkl.element_type, + num_bits_per_copy=mSFA_mkl.element_type.width, + ) + tiled_copy_sfa = cute.make_tiled_copy_tv( + atom_copy, cute.make_layout((32,)), cute.make_layout((1,)) + ) + tiled_copy_sfb = cute.make_tiled_copy_tv( + atom_copy, cute.make_layout((32,)), cute.make_layout((1,)) + ) + thr_copy_sfa = tiled_copy_sfa.get_slice(lane_idx) + thr_copy_sfb = tiled_copy_sfb.get_slice(lane_idx) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tAgSFA_mkl = thr_copy_sfa.partition_S(gSFA_mkl) + tAsSFA = thr_copy_sfa.partition_D(sSFA) + tAcSFA = thr_copy_sfa.partition_S(cSFA) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopN, loopK, loopL) + tBgSFB_nkl = thr_copy_sfb.partition_S(gSFB_nkl) + tBsSFB = thr_copy_sfb.partition_D(sSFB) + tBcSFB = thr_copy_sfb.partition_S(cSFB) + + # + # Partition shared/tensor memory tensor for TiledMMA_A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + tCrA = tiled_mma.make_fragment_A(sA) + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + + # + # Cluster wait before tensor memory alloc + # + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # + # Specialized Schedule warp + # + if warp_idx == self.sched_warp_id: + cute.arch.warpgroup_reg_dealloc(self.num_regs_sched_warps) + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + # First tile + work_tile = tile_sched.initial_work_tile_info() + cur_expert = self._tile_expert(work_tile, mTileExpert) + while work_tile.is_valid_tile & (cur_expert < 0): + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + cur_expert = self._tile_expert(work_tile, mTileExpert) + + tile_info_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_tile_stage + ) + + while work_tile.is_valid_tile: + # query next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + cur_expert = self._tile_expert(work_tile, mTileExpert) + while work_tile.is_valid_tile & (cur_expert < 0): + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + cur_expert = self._tile_expert(work_tile, mTileExpert) + + # acquire tile info pipeline + tile_info_pipeline.producer_acquire(tile_info_producer_state) + + # store the tile info + cur_tile_coord = work_tile.tile_idx + with cute.arch.elect_one(): + sInfo[(0, tile_info_producer_state.index)] = cur_tile_coord[0] + sInfo[(1, tile_info_producer_state.index)] = cur_tile_coord[1] + sInfo[(2, tile_info_producer_state.index)] = cur_expert + sInfo[(3, tile_info_producer_state.index)] = cutlass.Int32( + work_tile.is_valid_tile + ) + + # fence view async shared + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + self.sched_sync_barrier.arrive_and_wait() + # commit tile info pipeline + tile_info_pipeline.producer_commit(tile_info_producer_state) + tile_info_producer_state.advance() + + tile_info_pipeline.producer_tail(tile_info_producer_state) + + # + # Specialized TMA load warp + # + if warp_idx == self.tma_warp_id: + cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps) + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + # First tile + work_tile = tile_sched.initial_work_tile_info() + cur_expert = self._tile_expert(work_tile, mTileExpert) + while work_tile.is_valid_tile & (cur_expert < 0): + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + cur_expert = self._tile_expert(work_tile, mTileExpert) + + ab_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_ab_stage + ) + + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) + + # get the first tile info + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + # initialize the tile info + tile_info[0] = cur_tile_coord[0] + tile_info[1] = cur_tile_coord[1] + tile_info[2] = cur_expert + tile_info[3] = work_tile.is_valid_tile + + is_valid_tile = cutlass.Boolean(1) + is_valid_tile = tile_info[3] == 1 + + while is_valid_tile: + mma_tile_coord_mnl = ( + tile_info[0] // cute.size(tiled_mma.thr_id.shape), + tile_info[1], + tile_info[2], + ) + # + # Slice to per mma tile index + # + tAgA_slice = None + a_rows = None + if cutlass.const_expr(self.gather_a): + # Lane l gathers rows 4 l .. 4 l + 3 of the tile. + a_rows = cute.make_rmem_tensor((4,), cutlass.Int32) + a_row_base = tile_info[0] * self.cta_tile_shape_mnk[0] + lane_idx * 4 + for j in cutlass.range_constexpr(4): + a_rows[j] = mRowIndex[a_row_base + j] + else: + # ((atom_v, rest_v), loopK) + tAgA_slice = tAgA[(None, mma_tile_coord_mnl[0], None, 0)] + # ((atom_v, rest_v), loopK) + tBgB_slice = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])] + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + ab_producer_state.reset_count() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) + # + # Tma load loop + # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + a_stage = ab_producer_state.index + tAgA_k = None + tAsA_pipe = None + if cutlass.const_expr(not self.gather_a): + tAgA_k = tAgA_slice[(None, ab_producer_state.count)] + tAsA_pipe = tAsA[(None, ab_producer_state.index)] + tBgB_k = tBgB_slice[(None, ab_producer_state.count)] + tBsB_pipe = tBsB[(None, ab_producer_state.index)] + + tma_bar = ab_pipeline.producer_get_barrier(ab_producer_state) + + # Conditionally wait for AB buffer empty + ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status) + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 + ab_producer_state.advance() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) + + # TMA load A/B + if cutlass.const_expr(self.gather_a): + a_dst = sA.iterator + ( + a_stage * self.cta_tile_shape_mnk[0] * self.mma_tiler[2] + + lane_idx * 4 * self.mma_tiler[2] + ) + if cutlass.const_expr(self.is_a_mcast): + # The cluster's CTAs split the row groups and multicast them. + if lane_idx % self.num_mcast_ctas_a == block_in_cluster_coord_vmnk[2]: + _tma_gather4_load( + tma_atom_a, + a_dst, + tma_bar, + k_tile * self.mma_tiler[2], + a_rows, + a_full_mcast_mask, + ) + else: + _tma_gather4_load( + tma_atom_a, + a_dst, + tma_bar, + k_tile * self.mma_tiler[2], + a_rows, + None, + ) + else: + cute.copy( + tma_atom_a, + tAgA_k, + tAsA_pipe, + tma_bar_ptr=tma_bar, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_k, + tBsB_pipe, + tma_bar_ptr=tma_bar, + mcast_mask=b_full_mcast_mask, + ) + + # + # Advance to next tile + # + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + # + # Wait A/B buffer empty + # + ab_pipeline.producer_tail(ab_producer_state) + + # + # Specialized Scale load warp + # + if warp_idx == self.scale_warp_id: + cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps) + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + # First tile + work_tile = tile_sched.initial_work_tile_info() + cur_expert = self._tile_expert(work_tile, mTileExpert) + while work_tile.is_valid_tile & (cur_expert < 0): + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + cur_expert = self._tile_expert(work_tile, mTileExpert) + + scale_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_scale_stage + ) + + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) + + # get the first tile info + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + # initialize the tile info + tile_info[0] = cur_tile_coord[0] + tile_info[1] = cur_tile_coord[1] + tile_info[2] = cur_expert + tile_info[3] = work_tile.is_valid_tile + + is_valid_tile = cutlass.Boolean(1) + is_valid_tile = tile_info[3] == 1 + + while is_valid_tile: + # + # Prepare the mask for scaleA/scaleB + # + tApSFA = cute.make_rmem_tensor( + cute.make_layout( + cute.filter_zeros(cute.slice_(tAsSFA, (None, None, None, 0))).shape + ), + cutlass.Boolean, + ) + tBpSFB = cute.make_rmem_tensor( + cute.make_layout( + cute.filter_zeros(cute.slice_(tBsSFB, (None, None, None, 0))).shape + ), + cutlass.Boolean, + ) + sfa_rows = None + if cutlass.const_expr(self.gather_a): + # Lane l loads the scales of rows 4 l .. 4 l + 3, from their source rows. + # A padding row's index may be -1; any valid row serves for it. + sfa_rows = cute.make_rmem_tensor((4,), cutlass.Int32) + sfa_row_base = tile_info[0] * self.cta_tile_shape_mnk[0] + lane_idx * 4 + for j in cutlass.range_constexpr(4): + src_row = mRowIndex[sfa_row_base + j] + if src_row < 0: + src_row = cutlass.Int32(0) + sfa_rows[j] = src_row + + # Peek (try_wait) SCALE buffer empty + scale_producer_state.reset_count() + peek_scale_empty_status = cutlass.Boolean(1) + if scale_producer_state.count < k_tile_cnt: + peek_scale_empty_status = scale_pipeline.producer_try_acquire( + scale_producer_state + ) + + # + # load loop + # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # + # Slice to per mma tile index + # + tAsSFA_pipe = cute.filter_zeros( + tAsSFA[(None, None, None, scale_producer_state.index)] + ) + tBsSFB_pipe = cute.filter_zeros( + tBsSFB[(None, None, None, scale_producer_state.index)] + ) + tAgSFA_k = cute.filter_zeros( + tAgSFA_mkl[ + ( + None, + None, + None, + tile_info[0], + scale_producer_state.count, + 0, + ) + ] + ) + tBgSFB_k = cute.filter_zeros( + tBgSFB_nkl[ + ( + None, + None, + None, + tile_info[1], + scale_producer_state.count, + tile_info[2], + ) + ] + ) + + tAcSFA_compact = cute.filter_zeros( + cute.slice_( + tAcSFA, + ( + None, + None, + None, + tile_info[0], + scale_producer_state.count, + 0, + ), + ) + ) + tBcSFB_compact = cute.filter_zeros( + cute.slice_( + tBcSFB, + ( + None, + None, + None, + tile_info[1], + scale_producer_state.count, + tile_info[2], + ), + ) + ) + for i in cutlass.range_constexpr(cute.size(tApSFA, mode=[1])): + tApSFA[((0, 0), i, (0, 0))] = cute.elem_less( + tAcSFA_compact[(i)][0], mSFA_mkl.shape[0] + ) + for i in cutlass.range_constexpr(cute.size(tBpSFB, mode=[1])): + tBpSFB[((0, 0), i, (0, 0))] = cute.elem_less( + tBcSFB_compact[(i)][0], mSFB_nkl.shape[0] + ) + + # Conditionally wait for Scale buffer empty + scale_pipeline.producer_acquire(scale_producer_state, peek_scale_empty_status) + + # load scaleA/scaleB + if cutlass.const_expr(self.gather_a): + for j in cutlass.range_constexpr(4): + sfa_src = cute.make_tensor( + mSFX.iterator + + cute.crd2idx( + (sfa_rows[j], scale_producer_state.count, 0), mSFX.layout + ), + cute.make_layout((1,)), + ) + sfa_dst = cute.make_tensor( + sSFA.iterator + + ( + scale_producer_state.index * self.cta_tile_shape_mnk[0] + + lane_idx * 4 + + j + ), + cute.make_layout((1,)), + ) + cute.copy(atom_copy, sfa_src, sfa_dst) + else: + cute.copy(tiled_copy_sfa, tAgSFA_k, tAsSFA_pipe, pred=tApSFA) + cute.copy(tiled_copy_sfb, tBgSFB_k, tBsSFB_pipe, pred=tBpSFB) + + scale_pipeline.producer_commit(scale_producer_state) + + # Peek (try_wait) Scale buffer empty + scale_producer_state.advance() + peek_scale_empty_status = cutlass.Boolean(1) + if scale_producer_state.count < k_tile_cnt: + peek_scale_empty_status = scale_pipeline.producer_try_acquire( + scale_producer_state + ) + + # + # Advance to next tile + # + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + # + # Wait Scale buffer empty + # + scale_pipeline.producer_tail(scale_producer_state) + + # + # Specialized MMA warp + # + if warp_idx == self.mma_warp_id: + cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps) + # + # Bar sync for retrieve tensor memory ptr from shared mem + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + cur_expert = self._tile_expert(work_tile, mTileExpert) + while work_tile.is_valid_tile & (cur_expert < 0): + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + cur_expert = self._tile_expert(work_tile, mTileExpert) + + ab_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_ab_stage + ) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) + + # get the first tile info + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + # initialize the tile info + tile_info[0] = cur_tile_coord[0] + tile_info[1] = cur_tile_coord[1] + tile_info[2] = cur_expert + tile_info[3] = work_tile.is_valid_tile + + is_valid_tile = cutlass.Boolean(1) + is_valid_tile = tile_info[3] == 1 + + while is_valid_tile: + # Peek (try_wait) AB buffer full for k_tile = 0 + ab_consumer_state.reset_count() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) + + # Peek (try_wait) Acc buffer empty for k_tile = 0 + acc_producer_state.reset_count() + peek_acc_empty_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_acc_empty_status = acc_pipeline.producer_try_acquire(acc_producer_state) + + # + # Mma mainloop + # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Set tensor memory buffer for current tile + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + + # + # Wait for accumulator buffer empty + # + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state, peek_acc_empty_status) + + # + # Reset the ACCUMULATE field for each tile + # + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + if is_leader_cta: + # Conditionally wait for AB buffer full + ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status) + + # tCtAcc += tCrA * tCrB + num_kblocks = cute.size(tCrA, mode=[2]) + for kblock_idx in cutlass.range(num_kblocks, unroll_full=True): + kblock_coord = ( + None, + None, + kblock_idx, + ab_consumer_state.index, + ) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblock_coord], + tCrB[kblock_coord], + tCtAcc, + ) + # Enable accumulate on tCtAcc after first kblock + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + # Async arrive AB buffer empty + ab_pipeline.consumer_release(ab_consumer_state) + + # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 + ab_consumer_state.advance() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt: + if is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) + + # + # Async arrive accumulator buffer full(each kblock) + # + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + + # Peek (try_wait) Acc buffer empty for k_tile = k_tile + 1 + acc_producer_state.advance() + if acc_producer_state.count < k_tile_cnt: + if is_leader_cta: + peek_acc_empty_status = acc_pipeline.producer_try_acquire( + acc_producer_state + ) + + # + # Advance to next tile + # + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + # + # Wait for accumulator buffer empty + # + acc_pipeline.producer_tail(acc_producer_state) + + # + # Specialized acc update warps + # + if warp_idx <= self.acc_update_warp_id[-1]: + cute.arch.warpgroup_reg_alloc(self.num_regs_acc_update_warps) + # + # Bar sync for retrieve tensor memory ptr from shared memory + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + tCtAcc_final = cute.make_tensor( + tCtAcc_base.iterator + self.tmem_final_offset, tCtAcc_base.layout + ) + + # + # Partition for epilogue + # + epi_tidx = tidx % 128 + ( + tiled_copy_t2r, + tiled_copy_r2t, + tTR_tAcc_base, + tTR_rAcc, + tTR_rAcc_final, + tTR_sSFA, + tTR_sSFB, + tRT_rAcc, + tRT_tAcc_base, + ) = self.acc_update_tmem_copy_and_partition( + epi_tidx, + tCtAcc_base, + tCtAcc_final, + tCgC, + sSFA_view_as_C, + sSFB_view_as_C, + epi_tile, + ) + + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + cur_expert = self._tile_expert(work_tile, mTileExpert) + while work_tile.is_valid_tile & (cur_expert < 0): + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + cur_expert = self._tile_expert(work_tile, mTileExpert) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + + scale_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_scale_stage + ) + + epi_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1) + + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) + + # get the first tile info + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + # initialize the tile info + tile_info[0] = cur_tile_coord[0] + tile_info[1] = cur_tile_coord[1] + tile_info[2] = cur_expert + tile_info[3] = work_tile.is_valid_tile + + is_valid_tile = cutlass.Boolean(1) + is_valid_tile = tile_info[3] == 1 + + while is_valid_tile: + # initialize the final accumulator + tTR_rAcc_final.fill(0.0) + + tTR_rSFA = cute.make_rmem_tensor( + cute.slice_(tTR_sSFA, (None, None, None, 0, None, 0)).shape, + self.acc_dtype, + ) + tTR_rSFB = cute.make_rmem_tensor( + cute.slice_(tTR_sSFB, (None, None, None, 0, None, 0)).shape, + self.acc_dtype, + ) + + scale_consumer_state.reset_count() + peek_scale_full_status = cutlass.Boolean(1) + if scale_consumer_state.count < k_tile_cnt: + peek_scale_full_status = scale_pipeline.consumer_try_wait(scale_consumer_state) + + acc_consumer_state.reset_count() + peek_acc_full_status = cutlass.Boolean(1) + if acc_consumer_state.count < k_tile_cnt: + peek_acc_full_status = acc_pipeline.consumer_try_wait(acc_consumer_state) + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] + + # + # Wait for scale buffer full + # + scale_pipeline.consumer_wait(scale_consumer_state, peek_scale_full_status) + + tTR_sSFA_slice = cute.slice_( + tTR_sSFA, + (None, None, None, 0, None, scale_consumer_state.index), + ) + tTR_sSFB_slice = cute.slice_( + tTR_sSFB, + (None, None, None, 0, None, scale_consumer_state.index), + ) + + scale_atom_copy = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.acc_dtype, + num_bits_per_copy=self.acc_dtype.width, + ) + + cute.copy(scale_atom_copy, tTR_sSFA_slice, tTR_rSFA) + cute.copy(scale_atom_copy, tTR_sSFB_slice, tTR_rSFB) + + # + # Wait for accumulator buffer full + # + + acc_pipeline.consumer_wait(acc_consumer_state, peek_acc_full_status) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + + # + # Update accumulator by scale factor in subtiles + # + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + for subtile_idx in cutlass.range(subtile_cnt): + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + # + # Update accumulator by scale factor + # + tTR_rAcc_subtile = tTR_rAcc_final[(None, None, None, subtile_idx)] + tTR_rSFA_subtile = tTR_rSFA[(None, None, None, subtile_idx)] + tTR_rSFB_subtile = tTR_rSFB[(None, None, None, subtile_idx)] + + acc_vec = tTR_rAcc.load() + final_vec = tTR_rAcc_subtile.load() + scale_a = tTR_rSFA_subtile.load() + scale_b = tTR_rSFB_subtile.load() + scale = scale_a * scale_b + final_vec = acc_vec * scale + final_vec + tTR_rAcc_subtile.store(final_vec.to(self.acc_dtype)) + + # + # Async arrive accumulator buffer empty + # + scale_pipeline.consumer_release(scale_consumer_state) + scale_consumer_state.advance() + + peek_scale_full_status = cutlass.Boolean(1) + if scale_consumer_state.count < k_tile_cnt: + peek_scale_full_status = scale_pipeline.consumer_try_wait( + scale_consumer_state + ) + # + # Async arrive accumulator buffer empty + # + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + peek_acc_full_status = cutlass.Boolean(1) + if acc_consumer_state.count < k_tile_cnt: + peek_acc_full_status = acc_pipeline.consumer_try_wait(acc_consumer_state) + + tRT_tAcc = tRT_tAcc_base[(None, None, None, None, None, 0)] + tRT_tAcc = cute.group_modes(tRT_tAcc, 3, cute.rank(tRT_tAcc)) + + # + # Wait for epilogue buffer empty + # + epi_pipeline.producer_acquire(epi_producer_state) + + # copy the accumulator to tensor memory buffer + cute.copy(tiled_copy_r2t, tTR_rAcc_final, tRT_tAcc) + cute.arch.fence_view_async_tmem_store() + + # + # Async arrive epilogue buffer full + # + epi_pipeline.producer_commit(epi_producer_state) + epi_producer_state.advance() + + # + # Advance to next tile + # + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + # + # Specialized epilogue warps + # + if warp_idx <= self.epilog_warp_id[-1] and warp_idx >= self.epilog_warp_id[0]: + cute.arch.warpgroup_reg_alloc(self.num_regs_epilogue_warps) + # + # Alloc tensor memory buffer + # + tmem.allocate(self.num_tmem_alloc_cols) + + # + # Bar sync for retrieve tensor memory ptr from shared memory + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base_ = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + tCtAcc_final = cute.make_tensor( + tCtAcc_base_.iterator + self.tmem_final_offset, tCtAcc_base_.layout + ) + + # + # Partition for epilogue + # + epi_tidx = tidx % 128 + ( + tiled_copy_t2r, + tTR_tAcc_base, + tTR_rAcc, + ) = self.epilog_tmem_copy_and_partition( + epi_tidx, tCtAcc_final, tCgC, epi_tile, use_2cta_instrs + ) + + tTR_rC = None + tiled_copy_r2s = None + # simt_atom = None + tRS_rC = None + tRS_sC = None + bSG_sC = None + bSG_gC_partitioned = None + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition( + tiled_copy_t2r, tTR_rC, epi_tidx, sC + ) + ( + tma_atom_c, + bSG_sC, + bSG_gC_partitioned, + ) = self.epilog_gmem_copy_and_partition(epi_tidx, tma_atom_c, tCgC, epi_tile, sC) + + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + cur_expert = self._tile_expert(work_tile, mTileExpert) + while work_tile.is_valid_tile & (cur_expert < 0): + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + cur_expert = self._tile_expert(work_tile, mTileExpert) + + epi_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1) + + c_pipeline = None + # Threads/warps participating in tma store pipeline + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilog_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, + producer_group=c_producer_group, + ) + + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) + + # get the first tile info + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + # initialize the tile info + tile_info[0] = cur_tile_coord[0] + tile_info[1] = cur_tile_coord[1] + tile_info[2] = cur_expert + tile_info[3] = work_tile.is_valid_tile + + is_valid_tile = cutlass.Boolean(1) + is_valid_tile = tile_info[3] == 1 + + num_prev_subtiles = cutlass.Int32(0) + + while is_valid_tile: + mma_tile_coord_mnl = ( + tile_info[0] // cute.size(tiled_mma.thr_id.shape), + tile_info[1], + tile_info[2], + ) + # + # Slice to per mma tile index + # + bSG_gC = None + # ((ATOM_V, REST_V), EPI_M, EPI_N) + bSG_gC = bSG_gC_partitioned[ + ( + None, + None, + None, + mma_tile_coord_mnl[0], + mma_tile_coord_mnl[1], + 0, + ) + ] + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, epi_consumer_state.index)] + + # + # Wait for accumulator buffer full + # + epi_pipeline.consumer_wait(epi_consumer_state) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + + # + # Store accumulator to global memory in subtiles + # + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + for subtile_idx in cutlass.range(subtile_cnt): + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + # + # Convert to C type + # + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) + tRS_rC.store(acc_vec) + + # + # Store C to shared memory + # + num_prev_subtiles = num_prev_subtiles + 1 + c_buffer = num_prev_subtiles % self.num_c_stage + cute.copy( + tiled_copy_r2s, + tRS_rC, + tRS_sC[(None, None, None, c_buffer)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + self.epilog_sync_barrier.arrive_and_wait() + + # + # TMA store C to global memory + # + if warp_idx == self.epilog_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, subtile_idx)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + self.epilog_sync_barrier.arrive_and_wait() + + # + # Async arrive accumulator buffer empty + # + epi_pipeline.consumer_release(epi_consumer_state) + epi_consumer_state.advance() + + # + # Advance to next tile + # + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + # + # Wait for C store complete + # + c_pipeline.producer_tail() + + def __init__( + self, + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + raster_along_m: bool = False, + gather_a: bool = False, + ) -> None: + super().__init__(cutlass.Float32, use_2cta_instrs, mma_tiler_mn, cluster_shape_mn) + self.raster_along_m = raster_along_m + self.gather_a = gather_a + + @cute.jit + def _tile_expert(self, work_tile, mTileExpert: cute.Tensor) -> cutlass.Int32: + """The expert of the 128-row group that holds the work tile's CTA tile, or -1.""" + expert = cutlass.Int32(-1) + if work_tile.is_valid_tile: + expert = mTileExpert[ + (work_tile.tile_idx[0] * self.cta_tile_shape_mnk[0]) // GROUP_ROW_ALIGNMENT + ] + return expert + + def _compute_grid_grouped( + self, + c: cute.Tensor, + cta_tile_shape_mnk: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + max_active_clusters: cutlass.Constexpr, + ) -> Tuple[utils.PersistentTileSchedulerParams, Tuple[int, int, int]]: + """Persistent grid over every (M tile, N tile) of ``c``. + + With ``raster_along_m=False`` consecutive work indices walk N first, so the CTAs that + run at the same time share their A tiles in L2 and A is read from HBM about once. + """ + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, + (*cluster_shape_mn, 1), + swizzle_size=1, + raster_along_m=self.raster_along_m, + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + return tile_sched_params, grid + + @staticmethod + def can_implement_grouped( + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + n: int, + k: int, + gather_a: bool = False, + ) -> bool: + """Whether a tile configuration can run the grouped GEMM. + + Every CTA of a cluster must work on one expert, so a cluster spans one 128-row group: + the MMA tile is at most 128 rows and the cluster only grows along N. The gathered A + load (one tile::gather4 of four rows per lane of the TMA warp) needs a 1-CTA MMA + with 128-row tiles. + """ + if gather_a and (use_2cta_instrs or mma_tiler_mn[0] != GROUP_ROW_ALIGNMENT): + return False + if not Sm100BlockwiseGemmKernel.is_valid_mma_tiler_and_cluster_shape( + use_2cta_instrs, mma_tiler_mn, cluster_shape_mn + ): + return False + if mma_tiler_mn[0] > GROUP_ROW_ALIGNMENT: + return False + if cluster_shape_mn[0] != (2 if use_2cta_instrs else 1): + return False + n_tile = mma_tiler_mn[1] * cluster_shape_mn[1] + return n % n_tile == 0 and k % 128 == 0 + + @cute.jit + def wrapper_grouped( + self, + rows: cutlass.Int32, + n: cutlass.Int32, + k: cutlass.Int32, + num_experts: cutlass.Int32, + a_ptr: cute.Pointer, + b_ptr: cute.Pointer, + a_sf_ptr: cute.Pointer, + b_sf_ptr: cute.Pointer, + tile_expert_ptr: cute.Pointer, + c_tensor: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + ): + """Build the tensors from raw pointers and launch. + + ``c_tensor`` is (rows, N, 1) and stays a tensor so TVM FFI can find the stream. + """ + sf_k = (k + 127) // 128 + sf_n = (n + 127) // 128 + a_tensor = cute.make_tensor( + a_ptr, layout=cute.make_ordered_layout((rows, k, 1), order=(1, 0, 2)) + ) + b_tensor = cute.make_tensor( + b_ptr, layout=cute.make_ordered_layout((n, k, num_experts), order=(1, 0, 2)) + ) + # [K / 128, rows] in memory: rows are contiguous. + sfa_tensor = cute.make_tensor( + a_sf_ptr, layout=cute.make_ordered_layout((rows, sf_k, 1), order=(0, 1, 2)) + ) + # [E, N / 128, K / 128] in memory: K blocks are contiguous. + sfb_tensor = cute.make_tensor( + b_sf_ptr, + layout=cute.make_ordered_layout((sf_n, sf_k, num_experts), order=(1, 0, 2)), + ) + tile_expert = cute.make_tensor( + tile_expert_ptr, layout=cute.make_layout((rows // GROUP_ROW_ALIGNMENT,)) + ) + self( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + tile_expert, + tile_expert, # row_index: unused without gather_a + max_active_clusters, + stream, + ) + + @cute.jit + def wrapper_gathered( + self, + rows: cutlass.Int32, + num_tokens: cutlass.Int32, + n: cutlass.Int32, + k: cutlass.Int32, + num_experts: cutlass.Int32, + sfx_stride: cutlass.Int32, + x_ptr: cute.Pointer, + b_ptr: cute.Pointer, + x_sf_ptr: cute.Pointer, + b_sf_ptr: cute.Pointer, + tile_expert_ptr: cute.Pointer, + row_index_ptr: cute.Pointer, + c_tensor: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + ): + """``wrapper_grouped`` with A and its scales read through ``row_index``.""" + sf_k = (k + 127) // 128 + sf_n = (n + 127) // 128 + x_tensor = cute.make_tensor( + x_ptr, layout=cute.make_ordered_layout((num_tokens, k, 1), order=(1, 0, 2)) + ) + # [K / 128, T] in memory with row stride ``sfx_stride``: tokens are contiguous. + sfx_tensor = cute.make_tensor( + x_sf_ptr, + layout=cute.make_layout((num_tokens, sf_k, 1), stride=(1, sfx_stride, 0)), + ) + b_tensor = cute.make_tensor( + b_ptr, layout=cute.make_ordered_layout((n, k, num_experts), order=(1, 0, 2)) + ) + sfb_tensor = cute.make_tensor( + b_sf_ptr, + layout=cute.make_ordered_layout((sf_n, sf_k, num_experts), order=(1, 0, 2)), + ) + tile_expert = cute.make_tensor( + tile_expert_ptr, layout=cute.make_layout((rows // GROUP_ROW_ALIGNMENT,)) + ) + row_index = cute.make_tensor(row_index_ptr, layout=cute.make_layout((rows,))) + self( + x_tensor, + b_tensor, + c_tensor, + sfx_tensor, + sfb_tensor, + tile_expert, + row_index, + max_active_clusters, + stream, + ) + + +GROUP_ROW_ALIGNMENT = 128 + +# (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, raster_along_m) +GroupedGemmTactic = Tuple[bool, Tuple[int, int], Tuple[int, int], bool] + +# Best measured per (N, K) on B200 at 13312 tokens x top-8 over 288 experts. +DEFAULT_TACTICS: Dict[Tuple[int, int], GroupedGemmTactic] = { + (1024, 4096): (False, (128, 128), (1, 2), False), + (4096, 512): (False, (128, 128), (1, 2), False), +} +FALLBACK_TACTIC: GroupedGemmTactic = (False, (128, 128), (1, 1), False) +# The gathered A load keeps up with the MMA only when the A multicast of a wide cluster splits +# the row groups: (1, 8) covers FC1's eight N tiles. +DEFAULT_GATHERED_TACTICS: Dict[Tuple[int, int], GroupedGemmTactic] = { + (1024, 4096): (False, (128, 128), (1, 8), False), +} + +#: A compiled launcher: called with the wrapper's runtime arguments (sizes, raw pointers and the +#: C tensor); returns nothing. +CompiledGemm = Callable[..., None] + +_compiled_kernels: Dict[Tuple[int, int, GroupedGemmTactic, bool], CompiledGemm] = {} + + +_OPERAND_DTYPES = { + "a": torch.float8_e4m3fn, + "b": torch.float8_e4m3fn, + "sfa": torch.float32, + "sfb": torch.float32, + "tile_expert": torch.int32, + "row_index": torch.int32, + "out": torch.bfloat16, +} + + +def _check_operands(**operands: torch.Tensor) -> None: + """Raise unless each operand has its dtype and is contiguous: the kernels read raw pointers + with dense layouts, so another dtype or a strided view is read as wrong values.""" + for name, t in operands.items(): + if t.dtype != _OPERAND_DTYPES[name]: + raise ValueError(f"{name} must be {_OPERAND_DTYPES[name]}; got {t.dtype}") + if not t.is_contiguous(): + raise ValueError(f"{name} must be contiguous; got strides {t.stride()}") + + +def _c_view(out: torch.Tensor) -> torch.Tensor: + """``out`` [rows, N] as the kernel's (rows, N, 1) C tensor.""" + return out.view(1, *out.shape).permute(1, 2, 0) + + +def _compile( + n: int, k: int, tactic: GroupedGemmTactic, out: torch.Tensor, gather_a: bool = False +) -> CompiledGemm: + """The kernel for ``(n, k, tactic, gather_a)``, compiled on its first use. The row count and + the expert count are runtime arguments, so no other call compiles again.""" + use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, raster_along_m = tactic + key = (n, k, tactic, gather_a) + compiled = _compiled_kernels.get(key) + if compiled is not None: + return compiled + if not Sm100GroupedBlockwiseGemmKernel.can_implement_grouped( + use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, n, k, gather_a + ): + raise ValueError(f"grouped blockwise FP8 GEMM cannot run N={n}, K={k} with {tactic}") + gemm = Sm100GroupedBlockwiseGemmKernel( + use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, raster_along_m, gather_a + ) + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ) + + def ptr(dtype: type[cutlass.Numeric], align: int = 16) -> cute.Pointer: + return make_ptr(dtype, 16, cute.AddressSpace.gmem, assumed_align=align) + + c_fake = cute.runtime.from_dlpack(_c_view(out)).mark_layout_dynamic(leading_dim=1) + if gather_a: + compiled = cute.compile( + gemm.wrapper_gathered, + out.shape[0], + 1, + n, + k, + 1, + 1, + ptr(cutlass.Float8E4M3FN), + ptr(cutlass.Float8E4M3FN), + ptr(cutlass.Float32), + ptr(cutlass.Float32), + ptr(cutlass.Int32, 4), + ptr(cutlass.Int32, 16), + c_fake, + max_active_clusters=max_active_clusters, + stream=cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--opt-level 2 --enable-tvm-ffi", + ) + _compiled_kernels[key] = compiled + return compiled + compiled = cute.compile( + gemm.wrapper_grouped, + out.shape[0], + n, + k, + 1, + ptr(cutlass.Float8E4M3FN), + ptr(cutlass.Float8E4M3FN), + ptr(cutlass.Float32), + ptr(cutlass.Float32), + ptr(cutlass.Int32, 4), + c_fake, + max_active_clusters=max_active_clusters, + stream=cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--opt-level 2 --enable-tvm-ffi", + ) + _compiled_kernels[key] = compiled + return compiled + + +def grouped_blockwise_fp8_gemm( + a: torch.Tensor, + sfa: torch.Tensor, + b: torch.Tensor, + sfb: torch.Tensor, + tile_expert: torch.Tensor, + out: torch.Tensor, + tactic: Union[GroupedGemmTactic, None] = None, +) -> torch.Tensor: + """Write ``out[tile] = (sfa * a[tile]) @ (sfb[e] * b[e]).T`` for every used 128-row tile. + + Args: + a: ``[rows, K]`` FP8 E4M3, row-major. ``rows`` is a multiple of 128 and the rows are + grouped by expert in segments that start at multiples of 128. + sfa: ``[K / 128, rows]`` FP32, the 1 x 128 activation scales, K-block-major as + ``fp8_quantize_1x128`` returns them. + b: ``[E, N, K]`` FP8 E4M3, K-major per expert. + sfb: ``[E, N / 128, K / 128]`` FP32, the 128 x 128 weight block scales. + tile_expert: ``[rows / 128]`` Int32 on the GPU, the expert of each 128-row tile, or + -1 for an unused tile. Unused tiles get no MMA and no store, so their ``out`` rows + keep their old contents; rows of a used tile past its expert's segment hold + whatever the padding input produces. + out: ``[rows, N]`` BF16, written in place and returned. + tactic: ``(use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, raster_along_m)``; + defaults to the tuned choice for (N, K). + + N must be a multiple of 128 times the cluster's N size and K a multiple of 128. Every + operand is read through its raw pointer with the dense layout above, so each must be + contiguous. The call reads no device data on the host, so it is safe inside CUDA graph + capture. + """ + rows, k = a.shape + num_experts, n, k_b = b.shape + if k_b != k or rows % GROUP_ROW_ALIGNMENT != 0: + raise ValueError(f"bad shapes: a {tuple(a.shape)}, b {tuple(b.shape)}") + sfb_shape = (num_experts, n // 128, k // 128) + if tuple(sfa.shape) != (k // 128, rows) or tuple(sfb.shape) != sfb_shape: + raise ValueError(f"bad scale shapes: sfa {tuple(sfa.shape)}, sfb {tuple(sfb.shape)}") + if tile_expert.numel() != rows // GROUP_ROW_ALIGNMENT or tile_expert.dtype != torch.int32: + raise ValueError("tile_expert must be Int32 with rows / 128 entries") + if tuple(out.shape) != (rows, n) or out.dtype != torch.bfloat16: + raise ValueError(f"bad output: {tuple(out.shape)} {out.dtype}") + _check_operands(a=a, b=b, sfa=sfa, sfb=sfb, tile_expert=tile_expert, out=out) + if tactic is None: + tactic = DEFAULT_TACTICS.get((n, k), FALLBACK_TACTIC) + compiled = _compile(n, k, tactic, out) + compiled( + rows, + n, + k, + num_experts, + a.data_ptr(), + b.data_ptr(), + sfa.data_ptr(), + sfb.data_ptr(), + tile_expert.data_ptr(), + _c_view(out), + ) + return out + + +def grouped_blockwise_fp8_gemm_gathered( + x: torch.Tensor, + sfx: torch.Tensor, + row_index: torch.Tensor, + b: torch.Tensor, + sfb: torch.Tensor, + tile_expert: torch.Tensor, + out: torch.Tensor, + tactic: Union[GroupedGemmTactic, None] = None, +) -> torch.Tensor: + """``grouped_blockwise_fp8_gemm`` with A read as ``x[row_index]``, without the copy. + + Row ``r`` of the grouped operand is ``x[row_index[r]]`` with scales + ``sfx[:, row_index[r]]``. The TMA warp loads each A tile with ``tile::gather4`` (four + source rows per instruction) and the scale warp copies the matching scales, so the + result is bitwise the one of gathering ``x`` and ``sfx`` first. + + Args: + x: ``[T, K]`` FP8 E4M3, the un-permuted activations, row-major. + sfx: ``[K / 128, T]`` FP32, their 1 x 128 scales, K-block-major (tokens contiguous; + the row stride may exceed ``T``). + row_index: ``[rows]`` Int32, the source row in ``x`` of every grouped row. Padding + rows may hold -1 (or any valid row); their output is don't-care. + b, sfb, tile_expert, out: as in ``grouped_blockwise_fp8_gemm``. + tactic: as in ``grouped_blockwise_fp8_gemm``; must be a 1-CTA 128-row tile. + """ + num_tokens, k = x.shape + num_experts, n, k_b = b.shape + rows = out.shape[0] + if k_b != k or rows % GROUP_ROW_ALIGNMENT != 0: + raise ValueError(f"bad shapes: x {tuple(x.shape)}, b {tuple(b.shape)}, out {rows}") + # The kernel reads token t of K block j at ``j * sfx.stride(0) + t``: the tokens must be + # adjacent (a one-token scale's column stride is never read). + if ( + tuple(sfx.shape) != (k // 128, num_tokens) + or (num_tokens > 1 and sfx.stride(1) != 1) + or (k > 128 and sfx.stride(0) < num_tokens) + ): + raise ValueError(f"bad sfx: {tuple(sfx.shape)} strides {sfx.stride()}") + if tuple(sfb.shape) != (num_experts, n // 128, k // 128): + raise ValueError(f"bad sfb: {tuple(sfb.shape)}") + if row_index.numel() != rows or row_index.dtype != torch.int32: + raise ValueError("row_index must be Int32 with one entry per output row") + if tile_expert.numel() != rows // GROUP_ROW_ALIGNMENT or tile_expert.dtype != torch.int32: + raise ValueError("tile_expert must be Int32 with rows / 128 entries") + if tuple(out.shape) != (rows, n) or out.dtype != torch.bfloat16: + raise ValueError(f"bad output: {tuple(out.shape)} {out.dtype}") + if sfx.dtype != torch.float32: + raise ValueError(f"sfx must be FP32; got {sfx.dtype}") + _check_operands(a=x, b=b, sfb=sfb, tile_expert=tile_expert, row_index=row_index, out=out) + if tactic is None: + tactic = DEFAULT_GATHERED_TACTICS.get((n, k), FALLBACK_TACTIC) + compiled = _compile(n, k, tactic, out, gather_a=True) + compiled( + rows, + num_tokens, + n, + k, + num_experts, + sfx.stride(0), + x.data_ptr(), + b.data_ptr(), + sfx.data_ptr(), + sfb.data_ptr(), + tile_expert.data_ptr(), + row_index.data_ptr(), + _c_view(out), + ) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/moe/prefill_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/moe/prefill_moe.py new file mode 100644 index 000000000000..e93808aee036 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/moe/prefill_moe.py @@ -0,0 +1,423 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The routed experts of a prefill step on a grouped block-FP8 GEMM with the checkpoint's scales. + +The GEMMs are ``grouped_blockwise_fp8_gemm``: the persistent CuTe DSL blockwise kernel with +each 128-row tile mapped to its expert by a device table, FP32 1 x 128 activation and +128 x 128 weight scales as the checkpoint stores them. On these shapes it runs FC1 and FC2 +in about two thirds of the TRTLLM-Gen FP8 MoE's time. + +The launches, for ``T`` rows and top-8 routing: + +* ``topk`` (``route``): sigmoid scores in the TRTLLM-Gen routing kernel's form, FP32 correction + bias, top-8 (ties to the lower expert index; a NaN score ranks below every number), + normalization and the routed scaling factor, one program per row. +* The expert counts and their 128-aligned segment starts, in torch, with no host + synchronization. The buffers are sized for the worst case, ``8 T + 127 E`` rows. +* ``place``: takes each (row, slot) pair's destination row in its expert's segment with an + atomic increment and records the source row of every destination row. +* The tile table: the expert of every 128-row tile (``searchsorted`` over the segment ends), + -1 for tiles past an expert's rows; the GEMMs skip those. +* FC1: ``grouped_blockwise_fp8_gemm_gathered``, which loads each tile's rows from ``x`` + through the source-row table (TMA ``gather4``) instead of a permuted copy. +* ``act``: clamped SwiGLU of the ``[up; gate]`` product, then the 1 x 128 FP8 quantization + of the intermediate (scale ``amax / 448``). +* FC2: ``grouped_blockwise_fp8_gemm`` on the intermediate, already in segment order. +* ``combine``: each row's weighted sum of its eight slots, added to ``addend`` (the shared + expert's output) in FP32 and stored as BF16. + +The Triton kernels take the row count only through the pair count, which they do not +specialize on, so one compilation serves every step. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl +import triton.language.extra.libdevice as tldevice + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +from .fp8_block_scale_moe_runner import dequantize_block_weight +from .grouped_blockwise_gemm import grouped_blockwise_fp8_gemm, grouped_blockwise_fp8_gemm_gathered + +_TOPK = 8 +_ALIGN = 128 +_ACT_ROWS = 4 +_PLACE_BLOCK = 1024 + + +@triton.jit +def _topk_kernel( + logits_ptr, + bias_ptr, + ids_ptr, + wts_ptr, + scaling, + E: tl.constexpr, + EP: tl.constexpr, + TOPK: tl.constexpr, +): + row = tl.program_id(0) + e = tl.arange(0, EP) + valid = e < E + logit = tl.load(logits_ptr + row * E + e, mask=valid, other=0.0).to(tl.float32) + # The sigmoid as the TRTLLM-Gen routing kernel of the generation steps computes it + # (`sigmoid_accurate`), with libdevice's tanhf. + score = 0.5 * tldevice.tanh(0.5 * logit) + 0.5 + choice = score + tl.load(bias_ptr + e, mask=valid, other=0.0) + # A NaN choice ranks below every number; `avail` keeps a row of equal choices (all NaN, say) + # picking TOPK distinct experts, lowest index first. + choice = tl.where(choice == choice, choice, float("-inf")) + avail = valid + total = 0.0 + for i in tl.static_range(TOPK): + best = tl.max(tl.where(avail, choice, float("-inf")), axis=0) + idx = tl.min(tl.where(avail & (choice == best), e, EP), axis=0) + idx = tl.minimum(idx, E - 1) + w = tl.sum(tl.where(e == idx, score, 0.0), axis=0) + tl.store(ids_ptr + row * TOPK + i, idx) + tl.store(wts_ptr + row * TOPK + i, w) + total += w + avail = avail & (e != idx) + for i in tl.static_range(TOPK): + w = tl.load(wts_ptr + row * TOPK + i) + tl.store(wts_ptr + row * TOPK + i, w / total * scaling) + + +@triton.jit(do_not_specialize=["pairs"]) +def _place_kernel( + ids_ptr, + cursor_ptr, + dest_ptr, + row_index_ptr, + pairs, + TOPK: tl.constexpr, + BLOCK: tl.constexpr, +): + pair = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + valid = pair < pairs + expert = tl.load(ids_ptr + pair, mask=valid, other=0) + # Each pair's row in its expert's segment; the order within a segment is arbitrary. + dest = tl.atomic_add(cursor_ptr + expert, 1, mask=valid) + tl.store(dest_ptr + pair, dest, mask=valid) + tl.store(row_index_ptr + dest, pair // TOPK, mask=valid) + + +# `cap` is a multiple of 128 at every call, so its specialization never changes. +@triton.jit(do_not_specialize=["pairs"]) +def _act_kernel( + h_ptr, + dest_ptr, + out_ptr, + out_scale_ptr, + pairs, + cap, + limit, + INTER: tl.constexpr, + NB: tl.constexpr, + ROWS: tl.constexpr, +): + p = tl.program_id(0) * ROWS + tl.arange(0, ROWS) + valid = p < pairs + dest = tl.load(dest_ptr + p, mask=valid, other=0).to(tl.int64) + j = tl.arange(0, INTER) + src = h_ptr + dest[:, None] * (2 * INTER) + j[None, :] + up = tl.load(src, mask=valid[:, None], other=0.0).to(tl.float32) + gate = tl.load(src + INTER, mask=valid[:, None], other=0.0).to(tl.float32) + gate = tl.minimum(gate, limit) + up = tl.minimum(tl.maximum(up, -limit), limit) + act = (gate / (1.0 + tl.exp(-gate)) * up).reshape(ROWS, NB, 128) + scale = tl.maximum(tl.max(tl.abs(act), axis=2), 1e-10) / 448.0 + q = (act / scale[:, :, None]).to(tl.float8e4nv).reshape(ROWS, INTER) + tl.store(out_ptr + dest[:, None] * INTER + j[None, :], q, mask=valid[:, None]) + blk = tl.arange(0, NB) + tl.store(out_scale_ptr + blk[None, :] * cap + dest[:, None], scale, mask=valid[:, None]) + + +@triton.jit +def _combine_kernel( + y_ptr, + dest_ptr, + wts_ptr, + addend_ptr, + out_ptr, + H: tl.constexpr, + TOPK: tl.constexpr, + BH: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + n = tl.program_id(1) * BH + tl.arange(0, BH) + acc = tl.load(addend_ptr + row * H + n).to(tl.float32) + for k in tl.static_range(TOPK): + dest = tl.load(dest_ptr + row * TOPK + k).to(tl.int64) + w = tl.load(wts_ptr + row * TOPK + k) + acc += w * tl.load(y_ptr + dest * H + n).to(tl.float32) + tl.store(out_ptr + row * H + n, acc.to(tl.bfloat16)) + + +def route( + logits: torch.Tensor, bias: torch.Tensor, scaling: float +) -> tuple[torch.Tensor, torch.Tensor]: + """The routing ``prefill_moe`` applies: ``(ids int32 [T, 8], weights FP32 [T, 8])``. + + Per row: ``s = sigmoid(logits)``, computed as ``0.5 * tanh(0.5 * x) + 0.5`` like the + TRTLLM-Gen routing kernel the generation steps run, the 8 experts of the highest ``s + bias`` + in descending order (ties to the lower expert index, a NaN ranking below every number), and + each chosen expert's weight ``s_e / sum(s_chosen) * scaling``. A row whose choices are all + NaN picks experts 0 to 7 with NaN weights. + + Args: + logits: FP32 ``[T, E]``, contiguous, the router's outputs. + bias: FP32 ``[E]``, the correction bias. + scaling: the routed scaling factor. + """ + rows, num_experts = logits.shape + ids = torch.empty((rows, _TOPK), dtype=torch.int32, device=logits.device) + wts = torch.empty((rows, _TOPK), dtype=torch.float32, device=logits.device) + _topk_kernel[(rows,)]( + logits, + bias, + ids, + wts, + scaling, + num_experts, + triton.next_power_of_2(num_experts), + _TOPK, + num_warps=4, + ) + return ids, wts + + +def route_reference( + logits: torch.Tensor, bias: torch.Tensor, scaling: float +) -> tuple[torch.Tensor, torch.Tensor]: + """``route`` in plain torch (``torch.sigmoid``): a stable descending sort of the choices, NaN + mapped to -inf.""" + score = torch.sigmoid(logits.float()) + choice = score + bias.float() + choice = torch.where(torch.isnan(choice), float("-inf"), choice) + ids = torch.sort(choice, dim=-1, descending=True, stable=True).indices[:, :_TOPK] + wts = score.gather(1, ids) + wts = wts / wts.sum(-1, keepdim=True) * scaling + return ids.to(torch.int32), wts + + +class PrefillMoeWorkspace: + """The scratch of :data:`prefill_moe`, allocated once for the largest step seen and sliced after. + + The scratch is about 0.13 MB per routed pair (2.2 GB at 16K tokens). Allocating it per + layer lets the caching allocator free and re-map device memory in the middle of a + step, and that synchronizes the device while the collectives of other ranks wait for + this rank, which can stall every rank. The first prefill (the warmup's + largest step) sizes the buffers, so the KV-cache budget accounts for them. The model owns one + and passes it to every call. + """ + + def __init__(self) -> None: + self.cap = 0 + self.buffers: dict[str, torch.Tensor] = {} + + def get(self, cap: int, h: int, inter: int, device: torch.device) -> dict[str, torch.Tensor]: + """Dense views of ``cap`` rows: ``row_index`` int32 ``[cap]``, ``h1`` BF16 ``[cap, + 2 * inter]``, ``act`` E4M3 ``[cap, inter]``, ``act_scale`` FP32 ``[inter / 128, cap]`` + (contiguous) and ``y`` BF16 ``[cap, h]``.""" + if cap > self.cap or self.buffers.get("y") is None or self.buffers["y"].device != device: + self.buffers = {} + self.buffers = dict( + row_index=torch.empty(cap, dtype=torch.int32, device=device), + h1=torch.empty((cap, 2 * inter), dtype=torch.bfloat16, device=device), + act=torch.empty((cap, inter), dtype=torch.float8_e4m3fn, device=device), + act_scale=torch.empty((inter // 128, cap), dtype=torch.float32, device=device), + y=torch.empty((cap, h), dtype=torch.bfloat16, device=device), + ) + self.cap = cap + b = self.buffers + # The scales are K-block-major with a row stride of `cap`, the layout the activation + # kernel writes and the FC2 GEMM reads: a prefix of the buffer, not a column slice. + scale_blocks = b["act_scale"].shape[0] + return dict( + row_index=b["row_index"][:cap], + h1=b["h1"][:cap], + act=b["act"][:cap], + act_scale=b["act_scale"].view(-1)[: scale_blocks * cap].view(scale_blocks, cap), + y=b["y"][:cap], + ) + + +class PrefillMoe(OpWrapper): + """``addend + sum_i w_i * expert_{e_i}(x)`` for every row of a prefill step. + + ``x`` is the FP8 ``[T, H]`` row-wise 1 x 128 quantization of the rows with scales + ``x_scale`` (``[H/128, T]`` FP32, K-block-major as ``fp8_quantize_1x128`` returns them). + ``logits`` are the router's FP32 ``[T, E]`` outputs, ``bias`` the FP32 ``[E]`` correction + bias; ``route`` gives the experts ``e_i`` and the weights ``w_i``, ``scaling`` included. + ``limit`` is the SwiGLU clamp. ``w31`` / ``s31`` are the routed experts' ``[E, 2I, H]`` FP8 + ``[up; gate]`` weights and their ``[E, 2I/128, H/128]`` scales, ``w2`` / ``s2`` the ``[E, H, + I]`` down weights and scales. ``addend`` is BF16 ``[T, H]``. ``workspace`` is the caller's + :class:`PrefillMoeWorkspace`. Returns BF16 ``[T, H]``. + """ + + ARCHS = frozenset({Arch.SM_100}) + + # GLM-5.3-Flash at moe_tp4: hidden 4096, intermediate 2048 / 4 = 512, 288 experts. The + # target runs every row of a step with context rows here (the rows gathered from all + # ranks), whatever the step's length. + CELLS: tuple[Cell, ...] = ( + Cell( + why="a 13K-token prompt: every row of a sequence-parallel prefill step", + spec=dict(t=13312, e=288, h=4096, i=512), + ), + Cell(why="a 37-token prompt: an odd row count", spec=dict(t=37, e=288, h=4096, i=512)), + Cell(why="a one-token prefill step", spec=dict(t=1, e=288, h=4096, i=512)), + ) + + note = """ + The FC1 output is BF16 before the SwiGLU and the intermediate is requantized 1 x 128, as in + the reference here. Top-8 ties go to the lower expert index. A row of NaN logits routes to + experts 0 to 7 with NaN weights, so its output is NaN; the other rows are unaffected. + """ + + def __call__( + self, + x: torch.Tensor, + x_scale: torch.Tensor, + logits: torch.Tensor, + bias: torch.Tensor, + scaling: float, + limit: float, + w31: torch.Tensor, + s31: torch.Tensor, + w2: torch.Tensor, + s2: torch.Tensor, + addend: torch.Tensor, + workspace: PrefillMoeWorkspace, + ) -> torch.Tensor: + rows, h = x.shape + num_experts = logits.shape[1] + inter = w2.shape[2] + dev = x.device + ids, wts = route(logits, bias, scaling) + flat = ids.view(-1) + counts = torch.zeros(num_experts, dtype=torch.int32, device=dev) + counts.scatter_add_(0, flat.long(), torch.ones_like(flat)) + padded = (counts + _ALIGN - 1) // _ALIGN * _ALIGN + starts = torch.cumsum(padded, 0, dtype=torch.int32) - padded + ends = starts + counts + cursor = starts.clone() + dest = torch.empty_like(flat) + + cap = (rows * _TOPK + num_experts * (_ALIGN - 1) + _ALIGN - 1) // _ALIGN * _ALIGN + ws = workspace.get(cap, h, inter, dev) + row_index = ws["row_index"].fill_(-1) + pairs = rows * _TOPK + _place_kernel[(triton.cdiv(pairs, _PLACE_BLOCK),)]( + flat, cursor, dest, row_index, pairs, _TOPK, _PLACE_BLOCK, num_warps=4 + ) + tile_row = torch.arange(0, cap, _ALIGN, device=dev, dtype=torch.int32) + tile_expert = torch.searchsorted(starts + padded, tile_row, right=True) + tile_expert = tile_expert.clamp_(max=num_experts - 1).to(torch.int32) + tile_expert = torch.where(tile_row < ends[tile_expert], tile_expert, -1) + h1 = ws["h1"] + grouped_blockwise_fp8_gemm_gathered(x, x_scale, row_index, w31, s31, tile_expert, h1) + act, act_scale = ws["act"], ws["act_scale"] + _act_kernel[(triton.cdiv(pairs, _ACT_ROWS),)]( + h1, dest, act, act_scale, pairs, cap, limit, inter, inter // 128, _ACT_ROWS, num_warps=4 + ) + y = ws["y"] + grouped_blockwise_fp8_gemm(act, act_scale, w2, s2, tile_expert, y) + out = torch.empty((rows, h), dtype=torch.bfloat16, device=dev) + _combine_kernel[(rows, h // 512)](y, dest, wts, addend, out, h, _TOPK, 512, num_warps=4) + return out + + def reference( + self, + x: torch.Tensor, + x_scale: torch.Tensor, + logits: torch.Tensor, + bias: torch.Tensor, + scaling: float, + limit: float, + w31: torch.Tensor, + s31: torch.Tensor, + w2: torch.Tensor, + s2: torch.Tensor, + addend: torch.Tensor, + workspace: PrefillMoeWorkspace, + ) -> torch.Tensor: + """``route_reference``, then each expert in FP32 on the dequantized operands with the + kernels' BF16 FC1 and FC2 outputs and 1 x 128 requantized intermediate. ``workspace`` + is not read.""" + + def quant_1x128(v: torch.Tensor) -> torch.Tensor: + blocks = v.view(v.shape[0], -1, 128) + amax = blocks.abs().amax(-1, keepdim=True).clamp(min=1e-10) + scale = amax / 448.0 + return ((blocks / scale).to(torch.float8_e4m3fn).float() * scale).view_as(v) + + rows = x.shape[0] + inter = w2.shape[2] + xf = (x.float().view(rows, -1, 128) * x_scale.float().t()[:, :, None]).view(rows, -1) + out = addend.float().clone() + ids, wts = route_reference(logits, bias, scaling) + ids = ids.long() + with true_fp32_matmul(): + for e in ids.unique().tolist(): + sel, slot = (ids == e).nonzero(as_tuple=True) + w1 = dequantize_block_weight(w31[e], s31[e]) + h1 = (xf[sel] @ w1.t()).to(torch.bfloat16).float() + u = h1[:, :inter].clamp(-limit, limit) + g = h1[:, inter:].clamp(max=limit) + act = quant_1x128(torch.nn.functional.silu(g) * u) + y = (act @ dequantize_block_weight(w2[e], s2[e]).t()).to(torch.bfloat16).float() + out.index_add_(0, sel, wts[sel, slot, None] * y) + return out.to(torch.bfloat16) + + def is_valid( + self, + x: torch.Tensor, + x_scale: torch.Tensor, + logits: torch.Tensor, + bias: torch.Tensor, + scaling: float, + limit: float, + w31: torch.Tensor, + s31: torch.Tensor, + w2: torch.Tensor, + s2: torch.Tensor, + addend: torch.Tensor, + workspace: PrefillMoeWorkspace, + ) -> None: + rows, h = x.shape + assert x.dtype == torch.float8_e4m3fn and x.is_contiguous() + assert x_scale.dtype == torch.float32 and x_scale.shape == (h // 128, rows) + assert x_scale.shape[1] == 1 or x_scale.stride(1) == 1 + assert logits.dtype == torch.float32 and logits.is_contiguous() and logits.shape[0] == rows + num_experts = logits.shape[1] + # Every expert id is read from the logits' width: an expert count that disagrees indexes + # the weights, the scales or the bias past their experts. + assert num_experts >= _TOPK, f"{num_experts} experts for top-{_TOPK}" + assert bias.dtype == torch.float32 and bias.shape == (num_experts,) + inter = w2.shape[2] + assert w31.shape == (num_experts, 2 * inter, h), f"w31 {tuple(w31.shape)}" + assert s31.shape == (num_experts, 2 * inter // 128, h // 128), f"s31 {tuple(s31.shape)}" + assert w2.shape == (num_experts, h, inter), f"w2 {tuple(w2.shape)}" + assert s2.shape == (num_experts, h // 128, inter // 128), f"s2 {tuple(s2.shape)}" + for t in (w31, w2): + assert t.dtype == torch.float8_e4m3fn and t.is_contiguous() + for t in (s31, s2): + assert t.dtype == torch.float32 and t.is_contiguous() + assert ( + addend.dtype == torch.bfloat16 and addend.shape == (rows, h) and addend.is_contiguous() + ) + assert h % 512 == 0 and inter % 128 == 0 + assert isinstance(workspace, PrefillMoeWorkspace) + + +prefill_moe = PrefillMoe() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/norm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/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/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/norm/rms_norm_gated_token_major.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/norm/rms_norm_gated_token_major.py new file mode 100644 index 000000000000..a48f15533ae1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/norm/rms_norm_gated_token_major.py @@ -0,0 +1,304 @@ +# Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/ops/triton/layernorm_gated.py +# Copyright (c) 2024, Tri Dao, Albert Gu. +# +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gated RMS norm of per-head rows, with the gate read from a token-major ``[tokens, heads, N]`` view. + +The kernel, its launch configuration and its PDL policy are those Modeling V1's +``trtllm::rms_norm_gated_token_major`` (``tensorrt_llm/_torch/modules/mamba/layernorm_gated.py``) uses for +these layouts; this entry does not import it. The V1 op's generic one-row-per-program kernel for other +layouts is not part of this entry, so it takes the multi-row kernel's layouts only. + +The entry's contract and reference are copied from kimi_k3_vl's +(``models/kimi_k3_vl/ops/norm/rms_norm_gated_token_major.py``), which calls the V1 op; the reference is +the same computation. qwen4_exp's entry (``models/qwen4_exp/ops/norm/rms_norm_gated_token_major.py``) +also calls the V1 op, launching the multi-row kernel itself for its prefill shapes. +""" + +from __future__ import annotations + +import functools +import os + +import torch +import triton +import triton.language as tl +import triton.language.extra.libdevice as tldevice + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +# Rows per program of the multi-row kernel. At the decode shape (thousands of 128-element rows) one row +# per CTA leaves the kernel launch-limited; 4 rows with 4 warps reproduces the single-row kernel's +# reduction order at about twice the throughput. +_MULTIROW_ROWS = 4 +_MULTIROW_NUM_WARPS = 4 +_MULTIROW_MAX_N = 256 + + +@functools.lru_cache(maxsize=None) +def _pdl_device_policy(device_index: int) -> tuple[bool, int]: + """Whether PDL is allowed on this device, and its SM count.""" + props = torch.cuda.get_device_properties(device_index) + sm_version = props.major * 10 + props.minor + allowed = sm_version >= 90 and os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + return allowed, props.multi_processor_count + + +def _multirow_pdl(grid_size: int, device_index: int) -> bool: + """Use PDL only when the launch grid cannot cover all SMs by itself.""" + allowed, num_sms = _pdl_device_policy(device_index) + return allowed and grid_size < num_sms + + +@triton.heuristics({"OUTPUT_FP8": lambda args: args["FP8_SCALE"] is not None}) +@triton.heuristics({"SAVE_RSTD": lambda args: args["Rstd"] is not None}) +@triton.jit +def _rms_norm_gated_fwd_multirow_kernel( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + Z, # pointer to the gate branch + Rstd, # optional pointer to the 1/std + FP8_SCALE, # static scale for an optional FP8 output + stride_x_row, + stride_y_row, + stride_z_tok, + M, # number of rows in X + eps, + N: tl.constexpr, # row length; power of two, whole row per program + ROWS: tl.constexpr, # rows per program + HEADS_PER_TOK: tl.constexpr, + GATE_SIGMOID: tl.constexpr, + WEIGHT_IS_DELTA: tl.constexpr, + SAVE_RSTD: tl.constexpr, + OUTPUT_FP8: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + EARLY_TRIGGER: tl.constexpr = False, +): + """Gated rmsnorm(x), several short rows per program. + + Z is addressed token-major: row r reads z at + (r // HEADS_PER_TOK) * stride_z_tok + (r % HEADS_PER_TOK) * N. With + HEADS_PER_TOK == heads it reads a [num_tokens, heads, N] view whose + (heads, N) block is contiguous per token, e.g. a column slice of a wider + projection. + """ + rows = tl.program_id(0) * ROWS + tl.arange(0, ROWS) + row_mask = rows < M + cols = tl.arange(0, N) + mask2d = row_mask[:, None] + if EARLY_TRIGGER: + # Dependents read Y only after their own griddepcontrol.wait. + tl.extra.cuda.gdc_launch_dependents() + # The weight is a parameter (not produced in-step), so it is fetched before + # the PDL wait; X (the producer's output) and Z (an earlier in-step kernel's + # output) are read only after it. + w = tl.load(W + cols).to(tl.float32) + if WEIGHT_IS_DELTA: + w += 1.0 + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + x_off = rows[:, None].to(tl.int64) * stride_x_row + cols[None, :] + x = tl.load(X + x_off, mask=mask2d, other=0.0).to(tl.float32) + var = tl.sum(x * x, axis=1) / N + rstd = 1.0 / tl.sqrt(var + eps) + if SAVE_RSTD: + tl.store(Rstd + rows, rstd, mask=row_mask) + y = x * rstd[:, None] * w[None, :] + tok = rows // HEADS_PER_TOK + head = rows % HEADS_PER_TOK + z_off = tok[:, None].to(tl.int64) * stride_z_tok + head[:, None] * N + cols[None, :] + z = tl.load(Z + z_off, mask=mask2d, other=0.0).to(tl.float32) + gate = tl.sigmoid(z) + if not GATE_SIGMOID: + gate *= z + y *= gate + if OUTPUT_FP8: + # RMSNorm first rounds to the input dtype, then static quantization + # multiplies by the rounded reciprocal of its input scale. + y = y.to(X.dtype.element_ty).to(tl.float32) + y *= tldevice.rcp_rn(tl.load(FP8_SCALE).to(tl.float32)) + y_off = rows[:, None].to(tl.int64) * stride_y_row + cols[None, :] + tl.store(Y + y_off, y.to(Y.dtype.element_ty), mask=mask2d) + if LAUNCH_WITH_PDL and not EARLY_TRIGGER: + # Release only after both Rstd (when present) and Y have been stored. + tl.extra.cuda.gdc_launch_dependents() + + +class RmsNormGatedTokenMajor(OpWrapper): + """Return ``rmsnorm(x) * weight * gate(z)``, one launch, a new tensor. + + ``gate`` is ``sigmoid(z)`` (``"sigmoid"``, KDA's output gate) or ``z * sigmoid(z)`` (``"silu"``, GDN's). + Row ``r = t * heads + h`` of ``x`` is gated by ``z[t, h, :]``. + + Args: + x: ``[tokens * heads, N]`` bf16 (or fp16 / fp32) rows with a unit column stride; the row stride is free. + z: ``[tokens, heads, N]`` gate whose ``(heads, N)`` block is dense per token; the token stride is free + (a column slice of a wider projection). + weight: ``[N]`` contiguous, any float dtype. + eps: added to the mean square. + fp8_scale: optional one-element fp32 tensor on x's device; with it the output is float8_e4m3fn. + gate_activation: ``"sigmoid"`` or ``"silu"``. + + Returns: + A new contiguous ``[tokens * heads, N]`` tensor in ``x.dtype`` (float8_e4m3fn with ``fp8_scale``). + """ + + ARCHS = frozenset({Arch.SM_100}) + # The op's own forms; the shapes the shipped target normalizes are in this entry's test. + CELLS: tuple[Cell, ...] = ( + Cell( + why="silu gate (z * sigmoid(z)): the other gate form, 16 heads of 128", + spec=dict(tokens=64, heads=16, width=128, eps=1e-5, gate_activation="silu"), + ), + *( + Cell( + why=f"row width {width}: another power-of-two row a program holds whole", + spec=dict(tokens=48, heads=4, width=width, eps=1e-6, gate_activation="sigmoid"), + ) + for width in (64, 256) + ), + Cell( + why="6 heads: a row count that is not a multiple of the 4-row program (masked tail rows)", + spec=dict(tokens=7, heads=6, width=128, eps=1e-5, gate_activation="sigmoid"), + ), + *( + Cell( + why=f"fp8 output at scale {scale}: rounded norm output, then static e4m3 quantization", + spec=dict( + tokens=64, + heads=16, + width=128, + eps=1e-5, + gate_activation="sigmoid", + fp8_scale=scale, + ), + ) + for scale in (0.5, 3.0) + ), + ) + note = """ + Norm, weight and gate only; projections are the caller's. Norm before gate: + the mean square is taken over x alone. fp32 math, one rounding to the output + dtype. Inputs are not mutated; the output is new. + + One kernel: whole rows per program, four rows per program, so the row width + must be a power of two of at most 256 and z's per-token (heads, N) block dense + (is_valid). The module op's generic one-row kernel for other layouts is not + part of this entry. + + fp8 quantization first rounds to x.dtype, then multiplies by the correctly + rounded reciprocal of the fp32 scale and rounds to e4m3; the kernel reads only + the first element of fp8_scale. + + PDL is used when TRTLLM_ENABLE_PDL is unset or 1, the device is SM90 or later + and the grid has fewer programs than the device has SMs: the weight is loaded + before the wait, dependents are released at the start. It changes + scheduling, not results. Triton specializes on the first call per width, + head count and dtype; warm it outside CUDA-graph capture. + """ + + def __call__( + self, + x: torch.Tensor, + z: torch.Tensor, + weight: torch.Tensor, + eps: float, + fp8_scale: torch.Tensor | None = None, + gate_activation: str = "silu", + ) -> torch.Tensor: + rows, width = x.shape + out = torch.empty_like(x, dtype=torch.float8_e4m3fn if fp8_scale is not None else x.dtype) + grid_size = triton.cdiv(rows, _MULTIROW_ROWS) + device_index = x.device.index + launch_with_pdl = _multirow_pdl(grid_size, device_index) + with torch.cuda.device(device_index): + _rms_norm_gated_fwd_multirow_kernel[(grid_size,)]( + x, + out, + weight, + z, + None, + fp8_scale, + x.stride(0), + out.stride(0), + z.stride(0), + rows, + eps, + N=width, + ROWS=_MULTIROW_ROWS, + HEADS_PER_TOK=z.shape[1], + GATE_SIGMOID=gate_activation == "sigmoid", + WEIGHT_IS_DELTA=False, + LAUNCH_WITH_PDL=launch_with_pdl, + num_warps=_MULTIROW_NUM_WARPS, + # Dependents (out_proj) prefetch their weights; they read Y only + # after their own griddepcontrol.wait. + EARLY_TRIGGER=_pdl_device_policy(device_index)[0], + launch_pdl=launch_with_pdl, + ) + return out + + def reference( + self, + x: torch.Tensor, + z: torch.Tensor, + weight: torch.Tensor, + eps: float, + fp8_scale: torch.Tensor | None = None, + gate_activation: str = "silu", + ) -> torch.Tensor: + """fp32: norm before gate, the gate read token-major, one rounding to the output dtype. With + ``fp8_scale``, the normalized value is rounded to ``x.dtype`` first, as the kernel does.""" + rows, width = x.shape + zf = z.float().reshape(rows, width) + xf = x.float() + y = xf / torch.sqrt(xf.pow(2).mean(-1, keepdim=True) + eps) * weight.float() + gate = torch.sigmoid(zf) + if gate_activation == "silu": + gate = gate * zf + y = y * gate + if fp8_scale is None: + return y.to(x.dtype) + return (y.to(x.dtype).float() * (1.0 / fp8_scale.float())).to(torch.float8_e4m3fn) + + def is_valid( + self, + x: torch.Tensor, + z: torch.Tensor, + weight: torch.Tensor, + eps: float, + fp8_scale: torch.Tensor | None = None, + gate_activation: str = "silu", + ) -> None: + # The kernel checks nothing: every condition below is one it computes wrongly (or out of bounds) + # rather than refuses. + assert gate_activation in ("silu", "sigmoid"), ( + f"gate_activation must be 'silu' or 'sigmoid'; {gate_activation!r} computes the silu gate" + ) + assert x.ndim == 2 and z.ndim == 3, "x must be [tokens * heads, N] and z [tokens, heads, N]" + rows, width = x.shape + assert z.shape[2] == width and z.shape[0] * z.shape[1] == rows, ( + f"z {tuple(z.shape)} must be [tokens, heads, {width}] with tokens * heads == {rows} x rows: " + "row r reads z[r // heads, r % heads]" + ) + assert width & (width - 1) == 0 and width <= _MULTIROW_MAX_N, ( + f"the row width must be a power of two of at most {_MULTIROW_MAX_N}, got {width}" + ) + assert x.stride(-1) == 1, ( + "x's columns must be contiguous: the kernel steps them by one element" + ) + assert z.stride(2) == 1 and z.stride(1) == width, ( + "z's (heads, N) block must be dense per token: the kernel reads head h at h * N" + ) + assert weight.shape == (width,) and weight.is_contiguous(), ( + f"weight must be contiguous [{width}]" + ) + if fp8_scale is not None: + assert fp8_scale.numel() == 1, "fp8_scale must contain exactly one element" + assert fp8_scale.dtype == torch.float32, "fp8_scale must have dtype torch.float32" + assert fp8_scale.device == x.device, "fp8_scale must be on x's device" + + +rms_norm_gated_token_major = RmsNormGatedTokenMajor() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/quantization/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/quantization/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/quantization/__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/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/quantization/fp8_quantize_1x128.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/quantization/fp8_quantize_1x128.py new file mode 100644 index 000000000000..254e55b3580d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/quantization/fp8_quantize_1x128.py @@ -0,0 +1,156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dynamic 1x128 block-FP8 quantization of BF16 rows: e4m3 codes and one FP32 scale per 128 +columns of a row, the scales returned K-block-major (`[ceil(K/128), M]`) on the sm_100 family. + +Copied from qwen4_exp's entry (`models/qwen4_exp/ops/quantization/fp8_quantize_1x128.py`), which +wraps the same op. The reference here diverges from that one in two ways: `block_scales` divides +tensor by tensor (torch evaluates `scalar / tensor` as a rounded reciprocal times the scalar, up +to one ulp off the quotient the kernel computes), and the reference clamps to +-448 before the +E4M3 cast, as the kernel's saturating conversion does. The cells are at GLM's widths. +""" + +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 + +#: Columns per scale. +BLOCK = 128 + +#: Largest finite e4m3 magnitude. +E4M3_MAX = 448.0 + +#: The kernel's amax floor, as the bf16 constant it compares against: +#: bfloat16(1e-10) == 1.000444171950221e-10. +AMAX_FLOOR = float(torch.tensor(1e-10, dtype=torch.bfloat16)) + + +def block_scales(input: torch.Tensor, use_ue8m0: bool) -> tuple[torch.Tensor, torch.Tensor]: + """`(quant_scale, dequant_scale)`, each `[M, ceil(K/128)]` fp32, exactly rounded. + + The amax runs in bf16, as the kernel's running maximum does (exact: `|x|` of a bf16 is a + bf16), floored at `AMAX_FLOOR`; a short tail block takes the amax of its own elements. The + kernel multiplies by the first before rounding to e4m3 and returns the second. + + Every quotient is a tensor-by-tensor division: torch evaluates `scalar / tensor` as a + reciprocal times the scalar, which is two roundings and can sit one ulp off the quotient. + """ + rows, cols = input.shape + nb = (cols + BLOCK - 1) // BLOCK + mag = torch.nn.functional.pad(input.abs(), (0, nb * BLOCK - cols)) + amax = mag.reshape(rows, nb, BLOCK).amax(dim=2).float().clamp(min=AMAX_FLOOR) + one = torch.ones_like(amax) + qs = torch.full_like(amax, E4M3_MAX) / amax + if not use_ue8m0: + return qs, one / qs + raw = one / qs + mant, exp = torch.frexp(raw) + ds = torch.ldexp(one, torch.where(mant == 0.5, exp - 1, exp)) + return one / ds, ds + + +#: GLM-5.3-Flash's hidden width, the input of every routed and shared FFN GEMM. +_GLM_HIDDEN = 4096 + + +class Fp8Quantize1x128(OpWrapper): + """Quantize a bf16 `[M, K]` matrix to e4m3 with one fp32 scale per 1x128 row block. + + Per row `m` and block `b` of 128 consecutive columns (the last one short when + `K % 128 != 0`): `amax = max(|x[m, block]|, bf16(1e-10))`, `data = e4m3(x * (448 / amax))` + saturating at +-448, `scales[b, m] = 1 / (448 / amax)`. Returns `(data [M, K] + float8_e4m3fn, scales [ceil(K/128), M] fp32)`; `data * scales.t()` (each scale repeated + over its 128 columns) reconstructs `x`. + + The scales are TRANSPOSED, K-block-major: `[K/128, M]`, the layout + `cute_dsl_fp8_gemm_blackwell` and `fp8_block_scale_moe_runner` read on the sm_100 family. + Pass them on as returned. With `use_ue8m0=False` each scale is the plain `amax / 448`; + `use_ue8m0=True` rounds every scale up to a power of two -- a different quantization of + the same input, not a layout choice. + + Fusion boundary: the amax reduction, scaling, rounding and the compaction of the scales + into `[K/128, M]` are inside; no GEMM is fused. + """ + + ARCHS = frozenset({Arch.SM_100}) + + # Only GLM-5.3-Flash calls this entry, so its own cells are the op's row-count edges at + # GLM's hidden width; the widths and row counts the target quantizes are in the test. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one row: the decode row count, at GLM's hidden width", + spec=dict(rows=1, cols=_GLM_HIDDEN, use_ue8m0=False), + ), + Cell( + why="5 rows: not a multiple of 4, so data is a view of a padded allocation and the " + "scales a compacting copy", + spec=dict(rows=5, cols=_GLM_HIDDEN, use_ue8m0=False), + ), + Cell( + why="32 rows == K/128: the scale buffer is square, so its transpose is undetectable", + spec=dict(rows=32, cols=_GLM_HIDDEN, use_ue8m0=False), + ), + ) + + note = """ + Metadata consumed: none. The returned scale layout is chosen by an architecture query + (`isSM100Family()` in thop/fp8Quantize.cpp): off the sm_100 family the op returns a 1-D + padded buffer instead, which is why this entry is certified on sm_100 only. + + Preconditions, enforced loudly by the op: `input` bf16, 2-D (flatten `[B, S, K]` yourself), + contiguous, CUDA, and `K % 16 == 0` -- the width, not the block count. A consumer needing + `K % 128 == 0` (the block-scale MoE runner, the blockwise GEMM at every GLM shape) owns + that constraint. + + `data` is a leading-row view of an allocation padded to a multiple of 4 rows; the padded + rows are uninitialized, so a raw-pointer reader must stop at `M`. At a row count that is a + multiple of 4 nothing is padded and `scales` is the kernel's own buffer; otherwise it is a + compacting copy. + + At `M == K/128` the scales are square (32 rows at K 4096, 24 at 3072, 12 at 1536, 4 at + 512), so a consumer that reads them token-major passes every shape check and is silently + wrong. Nothing here can see that; the consumers' entries say so too. + + Source: compiled-in C++, `torch.ops.trtllm.fp8_quantize_1x128` (thop/fp8Quantize.cpp: + `fp8CS1x128` -> `launch_scale_1x128`, which takes `scale_1x128_vec_kernel` when the input + is 16-byte aligned and K % 8 == 0, as at every cell). The kernel's two fp32 divisions + (`448 / amax`, then `1 / that`) are IEEE round-to-nearest in the default build (`FAST_MATH` + off), the quotients `block_scales` computes, and its fp32 product and saturating E4M3 + conversion round as the reference's do. A `FAST_MATH=ON` build compiles the divisions + approximately (within 2 ulp each); the test's bands admit that build too. + + The target also runs this op inside its captured decode steps; the fp8_blockwise_gemm test + drives that capture (quantizer and GEMM at a tuned tactic). + + Not certified: `M == 0`, a short tail block (`K % 128 != 0`), all-zero blocks (the + `AMAX_FLOOR` scale), inf / NaN inputs, `use_ue8m0=True`, the element-wise fallback kernel + (an input whose storage offset breaks 16-byte alignment), sm_103. + """ + + def __call__( + self, + input: torch.Tensor, + use_ue8m0: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops.trtllm.fp8_quantize_1x128(input, use_ue8m0) + + def reference( + self, + input: torch.Tensor, + use_ue8m0: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + """`(data, scales[ceil(K/128), M])` from exactly rounded fp32 divisions, as a fresh + contiguous pair. + + No matmul, so nothing here depends on the TF32 setting. + """ + qs, ds = block_scales(input, use_ue8m0) + cols = qs.repeat_interleave(BLOCK, dim=1)[:, : input.shape[1]] + data = (input.float() * cols).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn) + return data, ds.t().contiguous() + + +fp8_quantize_1x128 = Fp8Quantize1x128() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/ssm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/ssm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/ssm/__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/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/ssm/kda_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/ssm/kda_decode.py new file mode 100644 index 000000000000..fdb6d0612657 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/ssm/kda_decode.py @@ -0,0 +1,271 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One-token KDA decode of B requests: the width-4 conv, the gated delta rule on each request's state row and +the gated output RMS norm, with the conv windows and state rows updated in place. + +Copied from kimi_k3_vl's entry (``models/kimi_k3_vl/ops/ssm/kda_decode.py``), which wraps the same +``trtllm::kda_decode`` op. This entry fixes the op's mode flags to the forms GLM-5.3-Flash runs (output norm, +lower-bound gate, beta sigmoid, conv windows updated in place); its reference computes every request at once and +returns the state and conv pools as the call leaves them, where kimi_k3_vl's loops over requests, takes every +mode flag and updates the pools in place. +""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 -- loads the library that registers trtllm::kda_decode +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: KDA's K = V. +_HEAD_DIM = 128 +#: The q / k L2-norm epsilon the kernel adds to the sum of squares. +_QK_NORM_EPS = 1e-6 + + +class KdaDecode(OpWrapper): + """Write the decode output of B requests into ``output``; update their conv windows and state rows. + + Per request ``b`` on row ``r`` (``ssm_state_indices[b]``, or ``b`` itself when it is None) and head ``h``: q, k + and v are ``silu(bias + sum_j w_t[j] * [window | x][j])`` over the three-tap window ``conv_state_*[r]`` + (oldest first) and the new input, which then replaces the oldest tap. q and k are L2-normalized with + ``rsqrt(sum(x^2) + 1e-6)`` (q then scaled); ``beta = sigmoid(beta_logit)``; the per-key decay is + ``exp(lower_bound * sigmoid(exp(a_log[h]) * (g + dt_bias)))``. With the V-first ``[V, K]`` state: + ``S <- S * decay``, ``S <- S + beta * (v - S k) k^T``, ``o = S q``, and the output is + ``o * rsqrt(mean(o^2) + onorm_eps) * onorm_weight * sigmoid(onorm_g)``. + + The op's mode flags are fixed to the forms the kernels implement: output norm, lower-bound gate and beta + sigmoid on, and the conv windows updated in place (``update_conv_cache``). + + Args: + x_q, x_k, x_v: bf16 ``[1, B, H, 128]`` new conv inputs; head and channel axes dense, the row stride free + (column sections of the fused projection). + w_q_t, w_k_t, w_v_t: contiguous bf16 ``[4, H * 128]`` conv taps (tap 3 multiplies the new input). + bias_q, bias_k, bias_v: contiguous bf16 ``[H * 128]`` conv biases (zeros for a bias-free conv). + conv_state_q, conv_state_k, conv_state_v: bf16 ``[rows, H * 128, 3]`` sections of one packed + ``[rows, 3 * H * 128, 3]`` conv pool (equal slot strides), updated in place. + a_log: contiguous fp32 ``[H]``. + g: bf16 ``[1, B, H, 128]`` raw forget-gate logits, row stride free. + dt_bias: contiguous fp32 ``[H * 128]``. + beta: bf16 ``[1, B, H]`` beta logits with a unit head stride. + onorm_g: bf16 ``[1, B, H, 128]`` output-gate logits, row stride free. + onorm_weight: contiguous fp32 ``[128]``. + ssm_state_indices: int32 ``[B]`` unique rows of the pools, or None for rows ``0..B-1``. + state: fp32 ``[rows, H, 128, 128]`` V-first state pool, updated in place; rows dense, slot stride a + multiple of 4 elements, 16-byte aligned. + lower_bound: the gate bound; scale: the q scale (``128 ** -0.5``); onorm_eps: the output norm epsilon. + output: contiguous bf16 ``[B, 1, H, 128]``, written. + + Returns: + ``output``. + """ + + ARCHS = frozenset({Arch.SM_100}) + # GLM-5.3-Flash's KDA at TP4 (16 of 64 heads). The op's own edges here; the decode steps the shipped target + # runs are in this entry's test. On sm_100 with 16 heads the kernel follows B * H: B <= 2 takes the four-CTA + # cluster kernel, 3 <= B <= 9 the legacy compact-heads one, 42 <= B <= 45 the legacy many-heads one, + # 39 <= B <= 54 otherwise and 86 <= B <= 156 the single-CTA cp.async one, other batches the two-stage + # bulk-copy one. + CELLS: tuple[Cell, ...] = ( + Cell( + why="two requests on scattered pool rows: the four-CTA cluster kernel's largest batch", + spec=dict(batch=2, layout="indexed", pool=12), + ), + Cell( + why="39 requests: the single-CTA cp.async kernel (624 <= B * H <= 865)", + spec=dict(batch=39, layout="indexed", pool=48), + ), + Cell( + why="44 requests: the legacy many-heads kernel (672 <= B * H <= 720)", + spec=dict(batch=44, layout="indexed", pool=48), + ), + ) + note = """ + Fused: conv, SiLU, q/k norm, gate and beta activation, the recurrence and the + gated RMS norm; the projections are the caller's. The op takes the conv + windows and states only by row: the batch metadata is ssm_state_indices. + With it, the rows are pool slots (a slot named twice races). Without it, + request b uses row b of state and of the conv sections, the dense form a + caller uses on fp32 staging copies of a bf16 pool's rows. + + The state must be fp32; the conv windows hold raw bf16 inputs, so their update + is a bit copy. All math is fp32 CUDA-core code with fast exp / rsqrt; the + output rounds once to bf16. + + One call per layer per step, in step order: swapping two steps of a request + computes another history without any error. CUDA-graph safe: nothing + allocates and nothing reads the device on the host. + + bias tensors are explicit: the op has no bias-free form, so a bias-free conv + passes zeros. The op rejects the other flag combinations and malformed layouts + itself (RuntimeError); is_valid states what it does not check. + """ + + def __call__( + self, + x_q: torch.Tensor, + x_k: torch.Tensor, + x_v: torch.Tensor, + w_q_t: torch.Tensor, + w_k_t: torch.Tensor, + w_v_t: torch.Tensor, + bias_q: torch.Tensor, + bias_k: torch.Tensor, + bias_v: torch.Tensor, + conv_state_q: torch.Tensor, + conv_state_k: torch.Tensor, + conv_state_v: torch.Tensor, + a_log: torch.Tensor, + g: torch.Tensor, + dt_bias: torch.Tensor, + beta: torch.Tensor, + onorm_g: torch.Tensor, + onorm_weight: torch.Tensor, + ssm_state_indices: torch.Tensor | None, + state: torch.Tensor, + lower_bound: float, + scale: float, + onorm_eps: float, + output: torch.Tensor, + ) -> torch.Tensor: + torch.ops.trtllm.kda_decode( + x_q, + x_k, + x_v, + w_q_t, + w_k_t, + w_v_t, + bias_q, + bias_k, + bias_v, + conv_state_q, + conv_state_k, + conv_state_v, + a_log, + g, + dt_bias, + beta, + onorm_g, + onorm_weight, + ssm_state_indices, + state, + True, # apply_onorm + True, # update_conv_cache + True, # use_lower_bound + True, # apply_beta_sigmoid + lower_bound, + scale, + onorm_eps, + output=output, + ) + return output + + def reference( + self, + x_q: torch.Tensor, + x_k: torch.Tensor, + x_v: torch.Tensor, + w_q_t: torch.Tensor, + w_k_t: torch.Tensor, + w_v_t: torch.Tensor, + bias_q: torch.Tensor, + bias_k: torch.Tensor, + bias_v: torch.Tensor, + conv_state_q: torch.Tensor, + conv_state_k: torch.Tensor, + conv_state_v: torch.Tensor, + a_log: torch.Tensor, + g: torch.Tensor, + dt_bias: torch.Tensor, + beta: torch.Tensor, + onorm_g: torch.Tensor, + onorm_weight: torch.Tensor, + ssm_state_indices: torch.Tensor | None, + state: torch.Tensor, + lower_bound: float, + scale: float, + onorm_eps: float, + output: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + """``(output, state, conv_state_q, conv_state_k, conv_state_v)`` as the call leaves them; mutates nothing. + + fp32 throughout (elementwise products and sums, no matmul); the output rounds once to bf16, the state + rows stay fp32 and the conv windows take the raw inputs. ``output``'s prior content is ignored. + """ + batch, heads = x_q.shape[1], x_q.shape[2] + if ssm_state_indices is None: + rows = torch.arange(batch, device=state.device) + else: + rows = ssm_state_indices.long() + acts, windows = [], [] + for x, w_t, bias, conv in ( + (x_q, w_q_t, bias_q, conv_state_q), + (x_k, w_k_t, bias_k, conv_state_k), + (x_v, w_v_t, bias_v, conv_state_v), + ): + full = torch.cat([conv[rows].float(), x[0].reshape(batch, -1, 1).float()], dim=-1) + act = (full * w_t.float().t()).sum(-1) + bias.float() + acts.append(torch.nn.functional.silu(act).reshape(batch, heads, _HEAD_DIM)) + windows.append(full[..., 1:]) + q, k, v = acts + q = q * torch.rsqrt(q.square().sum(-1, keepdim=True) + _QK_NORM_EPS) * scale + k = k * torch.rsqrt(k.square().sum(-1, keepdim=True) + _QK_NORM_EPS) + rate = a_log.float().exp().reshape(heads, 1) + gate = lower_bound * torch.sigmoid( + rate * (g[0].float() + dt_bias.float().reshape(heads, _HEAD_DIM)) + ) + b = beta[0].float().sigmoid() # [B, H] + s = state[rows].float() * gate.exp().unsqueeze(-2) # [B, H, V, K] + err = (v - (s * k.unsqueeze(-2)).sum(-1)) * b.unsqueeze(-1) + s = s + err.unsqueeze(-1) * k.unsqueeze(-2) + o = (s * q.unsqueeze(-2)).sum(-1) # [B, H, V] + o = o * torch.rsqrt(o.square().mean(-1, keepdim=True) + onorm_eps) + o = o * onorm_weight.float() * onorm_g[0].float().sigmoid() + new_state = state.clone() + new_state[rows] = s.to(state.dtype) + new_convs = [] + for conv, window in zip((conv_state_q, conv_state_k, conv_state_v), windows): + updated = conv.clone() + updated[rows] = window.to(conv.dtype) + new_convs.append(updated) + return (o.unsqueeze(1).to(output.dtype), new_state, *new_convs) + + def is_valid( + self, + x_q: torch.Tensor, + x_k: torch.Tensor, + x_v: torch.Tensor, + w_q_t: torch.Tensor, + w_k_t: torch.Tensor, + w_v_t: torch.Tensor, + bias_q: torch.Tensor, + bias_k: torch.Tensor, + bias_v: torch.Tensor, + conv_state_q: torch.Tensor, + conv_state_k: torch.Tensor, + conv_state_v: torch.Tensor, + a_log: torch.Tensor, + g: torch.Tensor, + dt_bias: torch.Tensor, + beta: torch.Tensor, + onorm_g: torch.Tensor, + onorm_weight: torch.Tensor, + ssm_state_indices: torch.Tensor | None, + state: torch.Tensor, + lower_bound: float, + scale: float, + onorm_eps: float, + output: torch.Tensor, + ) -> None: + # The op checks dtypes, shapes, strides and the state's alignment itself; it never reads the slots. + if ssm_state_indices is not None: + slots = ssm_state_indices.tolist() + assert len(set(slots)) == len(slots), ( + "ssm_state_indices must name each row once: two requests on a row race on its state" + ) + assert all(0 <= slot < state.shape[0] for slot in slots), ( + "ssm_state_indices must lie in the pool: the kernel indexes the state and conv rows unchecked" + ) + + +kda_decode = KdaDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/ssm/kda_mtp_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/ssm/kda_mtp_decode.py new file mode 100644 index 000000000000..49487381409c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/ssm/kda_mtp_decode.py @@ -0,0 +1,355 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""KDA speculative verify of B requests, 1 + num_spec new tokens each: replay of the previous round's accepted +drafts, then the conv, the gated delta rule and the gated output RMS norm of the new tokens, with the state and +conv window committed after the first (golden) token and the new drafts cached for the next round.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops # noqa: F401 -- registers trtllm::kda_mtp_decode +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: KDA's K = V, and the committed conv window (conv width 4 - 1). +_HEAD_DIM = 128 +_WINDOW = 3 +#: The q / k L2-norm epsilon the kernel adds to the sum of squares. +_QK_NORM_EPS = 1e-6 +#: The alignment the op promises its CuTe kernels for every tensor operand but the metadata and beta cache. +_ALIGN = 16 + + +class KdaMtpDecode(OpWrapper): + """Verify ``1 + num_spec`` packed new tokens per request in one launch; return the normalized output. + + Request ``n`` owns slot ``s = ssm_state_indices[n]`` and rows ``n * (1 + num_spec) + j`` of the token inputs. + Its slot holds the state after its last golden token, the raw conv inputs ``cs_*[s, :, 0:3]`` before it and, + from the previous round, ``num_spec`` drafts: their raw inputs ``cs_*[s, :, 3:3 + num_spec]`` and their + normalized k, log-decay g, v and beta in ``k_cache`` / ``g_cache`` / ``v_cache`` / ``beta_cache``. Of those, + the first ``a = num_accepted_tokens[s]`` were accepted. + + The kernel replays the a accepted drafts on the state (``S <- S * exp(g)``, ``S <- S + beta (v - S k) k^T`` + from the caches), then runs the new tokens: q, k, v are ``silu(sum_j w[:, j] * x[t + j])`` over the raw + sequence ``[window | accepted drafts | new tokens]``, q and k L2-normalized with ``rsqrt(sum(x^2) + 1e-6)`` + (q then scaled), ``beta = sigmoid(beta_logit)``, ``g = lower_bound * sigmoid(exp(A_log[h]) * (g_raw + + dt_bias))``, then the same recurrence and ``o = S q``. The output is + ``o * rsqrt(mean(o^2) + onorm_eps) * onorm_weight * sigmoid(onorm_g)``. + + In place, at slot s: the state after the golden token (rounded to the pool dtype), the conv window of the + three raw inputs ending at the golden token, and for draft ``j >= 1`` its raw inputs in conv column + ``2 + j`` and its k, g, v and beta in cache row ``j - 1``. ``num_accepted_tokens`` is read, not written: + the runtime records the next round's count after sampling. + + The op's mode is fixed to the runtime's verify: packed token rows (no ``cu_seqlens``), slot-indexed accepted + counts, the gated RMS norm fused, no metadata hints. + + Args: + x_q, x_k, x_v: bf16 ``[1, B * (1 + num_spec), H, 128]`` new conv inputs; head and channel axes dense, the + row stride free (column sections of the fused projection). + w_q, w_k, w_v: contiguous fp32 ``[H * 128, 4]`` conv taps. + cs_q, cs_k, cs_v: bf16 ``[slots, H * 128, >= 3 + num_spec]`` replay conv caches, dim-contiguous (stride 1 + along channels), updated in place. + g: bf16 ``[1, T, H, 128]`` raw forget-gate logits; beta: bf16 ``[1, T, H]`` beta logits, unit head stride. + A_log: fp32 ``[H]``; dt_bias: fp32 ``[H * 128]``. + recurrent_state: fp32 or bf16 ``[slots, H, 128, 128]`` V-first pool, updated in place. + k_cache, v_cache: bf16 (or fp32) ``[slots, num_spec, H * 128]``; g_cache: fp32 ``[slots, num_spec, H * 128]``; + beta_cache: bf16 (or fp32) ``[slots, num_spec, H]``; all updated in place. + ssm_state_indices: int32 ``[B]`` unique slots. + num_accepted_tokens: int32 ``[>= slots]``, slot-indexed (``prev_num_accepted_tokens``), each used entry in + ``[0, num_spec]``. + num_spec: drafts per round (at most 7). + lower_bound: the gate bound; scale: the q scale (``128 ** -0.5``). + onorm_g: bf16 ``[1, T, H, 128]`` output-gate logits; onorm_weight: fp32 ``[128]``; onorm_eps: its epsilon. + + Returns: + A new contiguous bf16 ``[1, B * (1 + num_spec), H, 128]`` output. + """ + + ARCHS = frozenset({Arch.SM_100}) + # GLM-5.3-Flash's KDA at TP4 (16 of 64 heads). The op's own edges here; the MTP-3 verify steps the shipped + # target runs are in this entry's test. With 16 heads the packed fused-norm verify runs the cluster-pair + # kernel while B * H <= 64, one CTA per (head, request) while B * H fits the SMs, the two-CTA-per-SM kernel + # beyond. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one draft per round, every request accepting it: the shortest replay, cluster-pair kernel", + spec=dict(batch=3, num_spec=1, pending=(1, 1, 1), state_dtype=torch.float32, pool=8), + ), + Cell( + why="seven drafts, accepted counts 0..7: the longest chunk the kernel holds (2 * 7 + 1 tokens)", + spec=dict( + batch=8, + num_spec=7, + pending=(0, 1, 2, 3, 4, 5, 6, 7), + state_dtype=torch.bfloat16, + pool=12, + ), + ), + ) + note = """ + Fused: conv, SiLU, q/k norm, gate and beta activation, replay, recurrence and + the gated RMS norm (on the bf16-rounded raw output, as the kernel's epilogue + does); projections are the caller's. The recurrence runs on tensor cores in + 3xTF32 (within about 2e-6 of fp32); gates, exponents and the norm are fp32 with + fast exp / rcp / rsqrt. The committed state is rounded to nearest for a bf16 + pool. Conv caches hold raw bf16 inputs and are bit copies. + + State protocol: after the call each used slot holds the state as of its golden + token, with the new drafts pending in the caches. The runtime records how many + were accepted in num_accepted_tokens before the next verify; nothing promotes + states on the host. The caches must be seeded from the live conv pool after a + prefill or a plain decode advanced it, and pending drafts committed before a + slot is released or moved. + + Each slot appears once per call. The kernel clamps a count above num_spec to + num_spec and does not guard a negative one (is_valid refuses both). + Every tensor operand is taken as 16-byte aligned except the int32 metadata + and the beta cache (whose alignment the op derives); the op does not check it. + Operands reach the kernels through DLPack, which refuses tensors that require + grad: pass parameters detached. The first call per layout compiles (CuTe + DSL); warm it outside CUDA-graph capture. The batch size is a runtime extent. + """ + + def __call__( + self, + x_q: torch.Tensor, + x_k: torch.Tensor, + x_v: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + k_cache: torch.Tensor, + g_cache: torch.Tensor, + v_cache: torch.Tensor, + beta_cache: torch.Tensor, + ssm_state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, + num_spec: int, + lower_bound: float, + scale: float, + onorm_g: torch.Tensor, + onorm_weight: torch.Tensor, + onorm_eps: float, + ) -> torch.Tensor: + return torch.ops.trtllm.kda_mtp_decode( + x_q=x_q, + x_k=x_k, + x_v=x_v, + w_q=w_q, + w_k=w_k, + w_v=w_v, + cs_q=cs_q, + cs_k=cs_k, + cs_v=cs_v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + k_cache=k_cache, + g_cache=g_cache, + v_cache=v_cache, + beta_cache=beta_cache, + ssm_state_indices=ssm_state_indices, + cu_seqlens=None, + num_spec=num_spec, + num_accepted_tokens=num_accepted_tokens, + lower_bound=lower_bound, + scale=scale, + zero_accepted_hint=False, + regular_metadata_hint=False, + packed_token_layout=True, + onorm_g=onorm_g, + onorm_weight=onorm_weight, + onorm_eps=onorm_eps, + fuse_output_norm=True, + ) + + def reference( + self, + x_q: torch.Tensor, + x_k: torch.Tensor, + x_v: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + k_cache: torch.Tensor, + g_cache: torch.Tensor, + v_cache: torch.Tensor, + beta_cache: torch.Tensor, + ssm_state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, + num_spec: int, + lower_bound: float, + scale: float, + onorm_g: torch.Tensor, + onorm_weight: torch.Tensor, + onorm_eps: float, + ) -> tuple[torch.Tensor, ...]: + """``(out, recurrent_state, cs_q, cs_k, cs_v, k_cache, g_cache, v_cache, beta_cache)`` as the call leaves + them; mutates nothing. + + Token by token in fp32 (elementwise products and sums, no matmul). Every store rounds once to its buffer's + dtype; ``out`` is normalized from the fp32 raw output (the kernel rounds that to bf16 first, a + difference of at most 2^-9 relative that the test's band covers). + """ + heads = A_log.shape[0] + steps = 1 + num_spec + rate = A_log.float().exp().reshape(heads, 1) + bias = dt_bias.float().reshape(heads, _HEAD_DIM) + weights = [w.float() for w in (w_q, w_k, w_v)] # [H * 128, 4] + tokens = [x[0].float().reshape(x.shape[1], -1) for x in (x_q, x_k, x_v)] # [T, H * 128] + out = torch.zeros(x_v.shape[1:], dtype=torch.float32, device=x_v.device) # [T, H, V] + state_after = recurrent_state.clone() + convs_after = [cs.clone() for cs in (cs_q, cs_k, cs_v)] + caches_after = [cache.clone() for cache in (k_cache, g_cache, v_cache, beta_cache)] + counts = num_accepted_tokens.tolist() + + def step( + state: torch.Tensor, + key: torch.Tensor, + log_decay: torch.Tensor, + value: torch.Tensor, + b: torch.Tensor, + ) -> torch.Tensor: + state = state * log_decay.exp().unsqueeze(-2) + err = (value - (state * key.unsqueeze(-2)).sum(-1)) * b.unsqueeze(-1) + return state + err.unsqueeze(-1) * key.unsqueeze(-2) + + for n, slot in enumerate(ssm_state_indices.tolist()): + accepted = min(counts[slot], num_spec) + state = recurrent_state[slot].float() # [H, V, K] + for t in range(accepted): + state = step( + state, + k_cache[slot, t].float().reshape(heads, _HEAD_DIM), + g_cache[slot, t].float().reshape(heads, _HEAD_DIM), + v_cache[slot, t].float().reshape(heads, _HEAD_DIM), + beta_cache[slot, t].float(), + ) + # The raw conv inputs in order: the committed window, the accepted drafts, the new tokens. + rows = slice(n * steps, (n + 1) * steps) + raw = [ + torch.cat([cs[slot, :, : _WINDOW + accepted].float().t(), x[rows]], dim=0) + for cs, x in zip((cs_q, cs_k, cs_v), tokens) + ] + for j in range(steps): + t = accepted + j + conv = [ + torch.nn.functional.silu( + sum(seq[t + tap] * w[:, tap] for tap in range(_WINDOW + 1)) + ).reshape(heads, _HEAD_DIM) + for seq, w in zip(raw, weights) + ] + q, key, value = conv + q = q * torch.rsqrt(q.square().sum(-1, keepdim=True) + _QK_NORM_EPS) * scale + key = key * torch.rsqrt(key.square().sum(-1, keepdim=True) + _QK_NORM_EPS) + row = n * steps + j + log_decay = lower_bound * torch.sigmoid(rate * (g[0, row].float() + bias)) + b = beta[0, row].float().sigmoid() + state = step(state, key, log_decay, value, b) + o = (state * q.unsqueeze(-2)).sum(-1) # [H, V] + o = o * torch.rsqrt(o.square().mean(-1, keepdim=True) + onorm_eps) + out[row] = o * onorm_weight.float() * onorm_g[0, row].float().sigmoid() + if j == 0: + state_after[slot] = state.to(state_after.dtype) + for cs_after, seq in zip(convs_after, raw): + cs_after[slot, :, :_WINDOW] = ( + seq[t + 1 : t + 1 + _WINDOW].t().to(cs_after.dtype) + ) + else: + for cs_after, seq in zip(convs_after, raw): + cs_after[slot, :, _WINDOW + j - 1] = seq[t + _WINDOW].to(cs_after.dtype) + for cache, value_row in zip(caches_after, (key, log_decay, value, b)): + cache[slot, j - 1] = value_row.reshape(cache.shape[2:]).to(cache.dtype) + return (out.unsqueeze(0).to(x_v.dtype), state_after, *convs_after, *caches_after) + + def is_valid( + self, + x_q: torch.Tensor, + x_k: torch.Tensor, + x_v: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + k_cache: torch.Tensor, + g_cache: torch.Tensor, + v_cache: torch.Tensor, + beta_cache: torch.Tensor, + ssm_state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, + num_spec: int, + lower_bound: float, + scale: float, + onorm_g: torch.Tensor, + onorm_weight: torch.Tensor, + onorm_eps: float, + ) -> None: + # The op checks shapes, dtypes, unit strides and the vectorized layouts itself; it never reads the + # metadata on the host and takes the operand alignment on trust. + slots = ssm_state_indices.tolist() + assert len(set(slots)) == len(slots), ( + "ssm_state_indices must name each slot once: two requests on a slot race on its state and caches" + ) + assert all(0 <= slot < recurrent_state.shape[0] for slot in slots), ( + "ssm_state_indices must lie in the pool: the kernel indexes every pool with them unchecked" + ) + # A count tensor shorter than the pool is the op's own (loud) refusal; only the values are checked here. + counts = num_accepted_tokens.tolist() + assert all(0 <= counts[slot] <= num_spec for slot in slots if slot < len(counts)), ( + f"num_accepted_tokens of a used slot must be in [0, {num_spec}]: a larger count is clamped and " + "a negative one shifts the request onto other requests' token rows" + ) + operands = dict( + x_q=x_q, + x_k=x_k, + x_v=x_v, + w_q=w_q, + w_k=w_k, + w_v=w_v, + cs_q=cs_q, + cs_k=cs_k, + cs_v=cs_v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + k_cache=k_cache, + g_cache=g_cache, + v_cache=v_cache, + onorm_g=onorm_g, + onorm_weight=onorm_weight, + ) + misaligned = [name for name, t in operands.items() if t.data_ptr() % _ALIGN] + assert not misaligned, ( + f"{misaligned} must start 16-byte aligned: the kernels assume it unchecked" + ) + + +kda_mtp_decode = KdaMtpDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/ssm/kda_prefill.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/ssm/kda_prefill.py new file mode 100644 index 000000000000..808d91cdd30b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/ssm/kda_prefill.py @@ -0,0 +1,243 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""KDA prefill: the gated delta rule over packed raw q / k / v from each sequence's resident state row, the +final state written back to the row. + +Copied from kimi_k3_vl's entry (``models/kimi_k3_vl/ops/ssm/kda_prefill.py``), which wraps the same +``trtllm::kda_prefill`` op. This entry adds the op's ``sm_reserve`` argument, and its reference returns the state +pool as the call leaves it (each final state rounded once to the pool dtype) where kimi_k3_vl's writes the pool in +place. +""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_custom_ops # noqa: F401 -- registers trtllm::kda_prefill +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: KDA's K = V. +_HEAD_DIM = 128 +#: Tokens per BT16 chunk; ``cu_chunks`` counts these. +_CHUNK = 16 + + +class KdaPrefill(OpWrapper): + """Run the gated delta rule over each sequence from ``state_pool[state_indices[b]]``; return ``o``. + + Per token, with q and k L2-normalized per head (q then scaled), ``beta = sigmoid(beta_logit)`` and the + per-key decay ``exp(lower_bound * sigmoid(exp(A_log[h]) * (g + dt_bias)))``: + ``S <- S * decay``, ``S <- S + beta * (v - S k) k^T``, ``o = S q``, with ``S`` the V-first ``[H, V, K]`` + state. The final state of sequence b is written back to its row, rounded to the pool dtype. + + Args: + q, k, v: contiguous bf16 ``[1, T, H, 128]`` raw (post-conv, unnormalized) inputs. + g: contiguous bf16 ``[1, T, H, 128]`` raw forget-gate logits (before ``dt_bias``). + beta: bf16 ``[1, T, H]`` beta logits; any token stride with a unit head stride. + state_pool: bf16 or fp32 ``[slots, H, 128, 128]`` V-first pool, updated in place; the slot stride is free + (16-byte multiple), rows dense. + state_indices: int32 ``[B]`` unique slots, one per sequence; a fresh sequence's row must be zeroed first. + cu_seqlens: int64 ``[B + 1]`` packed token offsets (``query_start_loc_long``). + A_log: contiguous fp32 ``[H]``; dt_bias: contiguous fp32 ``[H * 128]``. + scale: q scale (``128 ** -0.5``); lower_bound: the gate bound (finite, negative). + mode: ``"auto"``, ``"bt16"``, ``"engine"``, ``"decomp"`` or ``"small-bh"`` (the kernel route). + max_seqlen: zero, or a bound on every sequence length. + seq_order: optional int32 ``[B]`` stable descending-length order of the sequences. + cu_chunks: optional int32 ``[B + 1]`` prefix sums of ``ceil(length / 16)`` in request order; given + together with ``seq_order`` or not at all. + sm_reserve: SMs kept out of the Small-BH persistent grid. + + Returns: + A new contiguous bf16 ``[1, T, H, 128]`` output. + """ + + ARCHS = frozenset({Arch.SM_100}) + # GLM-5.3-Flash's KDA at TP4 (16 of 64 heads). The op's own edges; the prefill steps the shipped target runs + # (engine route, scheduler-planned metadata) are in this entry's test. + CELLS: tuple[Cell, ...] = ( + Cell( + why="auto route without planned metadata: on-device schedule, ragged lengths, empty sequence", + spec=dict( + heads=16, + seq_lens=(100, 0, 37, 300), + slots=(7, 2, 10, 4), + fresh=(True, False, True, False), + state_dtype=torch.float32, + mode="auto", + planned=False, + pool=12, + ), + ), + Cell( + why="engine route with planned metadata on a bf16 pool: chunk edges 15 / 16 / 17 / 63 / 65", + spec=dict( + heads=16, + seq_lens=(15, 16, 17, 63, 65), + slots=(0, 3, 5, 8, 11), + fresh=(True, False, True, False, True), + state_dtype=torch.bfloat16, + mode="engine", + planned=True, + pool=12, + ), + ), + ) + note = """ + Q/K normalization (x / max(|x|, 1e-12)), gate and beta activation are fused; + the convolution and the output norm are not. The kernel evaluates 16-token + chunks with bf16 tensor-core operands and fp32 accumulation, the state held in + fp32 across the call and rounded to the pool dtype once at the end. + + The caller zeroes fresh rows (reset_recurrent_state_rows) before the call: + every selected row is read as the initial state. Only selected rows are + written; other slots and slot padding are untouched. Each slot appears once. + + Routes: auto and bt16 pick the engine or the chunk decomposition by + occupancy (decomposition iff 2 * B * H fits the SMs); engine and decomp force + one. small-bh is opt-in only and needs a bf16 pool; with one it diverges on + long sequences mixing near-unit and strong decay. max_seqlen and sm_reserve + feed the Small-BH route only (and the argument checks). seq_order and + cu_chunks, when given, replace the schedule the kernel otherwise builds on + the device; they must be refreshed when the lengths change. None of the + scheduling arguments change the recurrence. + + The first call per (pool dtype, bound, route) compiles; warm it outside + CUDA-graph capture. Inputs must not alias the pool (the op checks). + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + state_pool: torch.Tensor, + state_indices: torch.Tensor, + cu_seqlens: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + scale: float, + lower_bound: float, + mode: str, + max_seqlen: int, + seq_order: torch.Tensor | None, + cu_chunks: torch.Tensor | None, + sm_reserve: int, + ) -> torch.Tensor: + return torch.ops.trtllm.kda_prefill( + q, + k, + v, + g, + beta, + state_pool, + state_indices, + cu_seqlens, + A_log, + dt_bias, + scale, + lower_bound, + mode, + max_seqlen, + seq_order, + cu_chunks, + sm_reserve, + ) + + def reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + state_pool: torch.Tensor, + state_indices: torch.Tensor, + cu_seqlens: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + scale: float, + lower_bound: float, + mode: str, + max_seqlen: int, + seq_order: torch.Tensor | None, + cu_chunks: torch.Tensor | None, + sm_reserve: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """``(o, state_pool as the call leaves it)``; mutates nothing. + + The token-by-token recurrence in fp32 (elementwise products and sums, no matmul), each final state + rounded once to the pool dtype and ``o`` once to bf16. ``mode``, ``max_seqlen``, ``seq_order``, + ``cu_chunks`` and ``sm_reserve`` schedule the kernel and are ignored. + """ + heads = q.shape[2] + qf = torch.nn.functional.normalize(q[0].float(), dim=-1, eps=1e-12) * scale + kf = torch.nn.functional.normalize(k[0].float(), dim=-1, eps=1e-12) + vf = v[0].float() + rate = A_log.float().exp().reshape(heads, 1) + gate = lower_bound * torch.sigmoid( + rate * (g[0].float() + dt_bias.float().reshape(heads, -1)) + ) + decay = gate.exp() # [T, H, K] + bf = beta[0].float().sigmoid() # [T, H] + out = torch.zeros_like(vf) + pool = state_pool.clone() + bounds = cu_seqlens.tolist() + for seq, slot in enumerate(state_indices.tolist()): + state = state_pool[slot].float() # [H, V, K] + for token in range(bounds[seq], bounds[seq + 1]): + state = state * decay[token].unsqueeze(-2) + key = kf[token].unsqueeze(-2) # [H, 1, K] + err = vf[token] - (state * key).sum(-1) # [H, V] + state = state + (bf[token].unsqueeze(-1) * err).unsqueeze(-1) * key + out[token] = (state * qf[token].unsqueeze(-2)).sum(-1) + pool[slot] = state.to(pool.dtype) + return out.unsqueeze(0).to(v.dtype), pool + + def is_valid( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + state_pool: torch.Tensor, + state_indices: torch.Tensor, + cu_seqlens: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + scale: float, + lower_bound: float, + mode: str, + max_seqlen: int, + seq_order: torch.Tensor | None, + cu_chunks: torch.Tensor | None, + sm_reserve: int, + ) -> None: + # The op checks dtypes, shapes, strides, alignment and aliasing; it never reads the metadata values on the + # host, and these are the ones it takes and gets wrong. + bounds = cu_seqlens.tolist() + lengths = [hi - lo for lo, hi in zip(bounds, bounds[1:])] + assert bounds[0] == 0 and bounds[-1] == q.shape[1] and min(lengths) >= 0, ( + "cu_seqlens must monotonically partition the packed tokens" + ) + slots = state_indices.tolist() + assert len(slots) == len(set(slots)), ( + "state_indices must name one unique slot per sequence: two sequences on a slot race" + ) + assert all(0 <= slot < state_pool.shape[0] for slot in slots), "state slot out of range" + assert max_seqlen == 0 or max(lengths) <= max_seqlen, ( + "max_seqlen is below a sequence length" + ) + if seq_order is not None and cu_chunks is not None: + assert seq_order.tolist() == sorted( + range(len(lengths)), key=lambda i: (-lengths[i], i) + ), "seq_order must be the stable descending sequence-length order" + chunks = [0] + for length in lengths: + chunks.append(chunks[-1] + (length + _CHUNK - 1) // _CHUNK) + assert cu_chunks.tolist() == chunks, "cu_chunks must prefix-sum ceil(length / 16)" + + +kda_prefill = KdaPrefill() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/ssm/kda_prefill_conv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/ssm/kda_prefill_conv.py new file mode 100644 index 000000000000..9011e0653188 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/ssm/kda_prefill_conv.py @@ -0,0 +1,184 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""KDA prefill short convolution: the causal width-4 conv + SiLU of the packed ``[q | k | v]`` projection, +each sequence's conv history read from and written back to its slot of the live conv pool.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_custom_ops # noqa: F401 -- registers trtllm::kda_prefill_conv +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Channels of one head (KDA's K = V); the op's q / k / v outputs are ``[1, T, H, 128]``. +_HEAD_DIM = 128 + + +class KdaPrefillConv(OpWrapper): + """``(q, k, v)`` = SiLU of the causal width-4 convolution of the packed projection, one launch. + + Sequence ``b`` is tokens ``starts[b]:starts[b + 1]`` of ``src``. Its raw inputs are preceded by a three-tap + history: ``conv_state[slots[b]]`` (oldest tap first) when ``has_init[b]``, zeros otherwise. For channel ``c`` + and token ``t``, ``y[t, c] = silu(sum_j weight[c, j] * seq[t + j, c])`` over that history-prefixed sequence, + tap 3 being token ``t`` itself. Channels ``[0, D)`` are q, ``[D, 2D)`` k and ``[2D, 3D)`` v (``D = H * 128``), + returned as ``[1, T, H, 128]`` with head ``(c % D) // 128``. The slot's history becomes the last three entries + of the history-prefixed sequence (raw bit copies), so a sequence shorter than three tokens keeps part of its + old history (a fresh one, part zeros). + + Args: + src: bf16 ``[T, >= 3D]`` token-major projection rows; columns past ``3D`` are ignored. + weight: bf16 ``[3D, 4]`` conv taps of q, k and v stacked (``conv1d.weight.squeeze(1)``). + conv_state: bf16 ``[slots, 3D, 3]`` live conv pool, updated in place; the slot stride is free. + starts: int32 ``[B + 1]`` packed token offsets (``query_start_loc``). + slots: int32 ``[B]`` unique pool slot of each sequence. + has_init: bool ``[B]``, whether the sequence continues from its slot's history. + max_seqlen: host int bound on the sequence lengths; it sizes the grid only (any value is correct). + + Returns: + ``(q, k, v)``, new contiguous bf16 ``[1, T, H, 128]`` tensors, unnormalized. + """ + + ARCHS = frozenset({Arch.SM_100}) + # GLM-5.3-Flash's KDA at TP4: 16 of 64 heads, 3D = 6144 channels. The op's own edges; the prefill steps the + # shipped target runs are in this entry's test. + CELLS: tuple[Cell, ...] = ( + Cell( + why="sequences of 1, 2, 3 and 5 tokens: the new history mixes the old window with new inputs", + spec=dict( + heads=16, + seq_lens=(1, 2, 3, 5), + slots=(6, 0, 9, 3), + has_init=(True, True, False, True), + pool=10, + ), + ), + Cell( + why="lengths off the 32-token run and 8-token group edges, fresh and continuing slots", + spec=dict( + heads=16, + seq_lens=(31, 33, 257, 100), + slots=(2, 7, 4, 5), + has_init=(False, True, True, False), + pool=8, + ), + ), + Cell( + why="a [q | k | v | g] source row wider than the conv channels (row stride 4D, g columns ignored)", + spec=dict( + heads=16, + seq_lens=(64, 17), + slots=(1, 0), + has_init=(True, False), + pool=4, + src_width=4 * 16 * _HEAD_DIM, + ), + ), + ) + note = """ + Fused: conv, SiLU and the q / k / v split, no normalization (the prefill + recurrence normalizes q and k). starts is int32 (query_start_loc); the + recurrence that follows takes the int64 copy (query_start_loc_long). + + SiLU: for x >= -4 the kernel uses x/2 + x/2 * tanh.approx(x/2) (relative error + at most about 2^-11 for x >= 0, absolute error under 1e-3 for -4 <= x < 0); + groups of tokens with any smaller x are recomputed with an exact form. Taps + accumulate in fp32 with exactly halved bf16 weights; one rounding to bf16. + + Every sequence of a call must name a distinct slot (two sequences on one slot + race on its history). Fresh sequences need no reset: has_init=False reads a + zero history and the write-back replaces the whole window. The kernel never + reads max_seqlen on the device; it sizes the grid, and CTAs loop over the + device-side run count, so an understated bound is still correct. + + The first call per layout compiles (CuTe DSL); warm it outside CUDA-graph + capture. Token count, batch size and pool capacity are runtime extents. + """ + + def __call__( + self, + src: torch.Tensor, + weight: torch.Tensor, + conv_state: torch.Tensor, + starts: torch.Tensor, + slots: torch.Tensor, + has_init: torch.Tensor, + max_seqlen: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return torch.ops.trtllm.kda_prefill_conv( + src, weight, conv_state, starts, slots, has_init, max_seqlen + ) + + def reference( + self, + src: torch.Tensor, + weight: torch.Tensor, + conv_state: torch.Tensor, + starts: torch.Tensor, + slots: torch.Tensor, + has_init: torch.Tensor, + max_seqlen: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """``(q, k, v, conv_state as the call leaves it)``; mutates nothing. + + Taps in fp32 in tap order and an exact fp32 SiLU, rounded to bf16 once. ``max_seqlen`` only sizes the + kernel's grid and is ignored. + """ + channels, width = weight.shape + history = width - 1 + heads = channels // (3 * _HEAD_DIM) + tokens = src.shape[0] + x = src[:, :channels].float() + w = weight.float() + y = torch.zeros(tokens, channels, dtype=torch.float32, device=src.device) + states = conv_state.clone() + bounds = starts.tolist() + for seq, (slot, carry) in enumerate(zip(slots.tolist(), has_init.tolist())): + lo, hi = bounds[seq], bounds[seq + 1] + length = hi - lo + if carry: + prefix = conv_state[slot].float().t() + else: + prefix = x.new_zeros(history, channels) + padded = torch.cat([prefix, x[lo:hi]], dim=0) + acc = torch.zeros(length, channels, dtype=torch.float32, device=src.device) + for tap in range(width): + acc = acc + padded[tap : tap + length] * w[:, tap] + y[lo:hi] = torch.nn.functional.silu(acc) + states[slot] = padded[length : length + history].t().to(states.dtype) + q, k, v = ( + part.reshape(1, tokens, heads, _HEAD_DIM).to(src.dtype) + for part in y.split(channels // 3, dim=1) + ) + return q, k, v, states + + def is_valid( + self, + src: torch.Tensor, + weight: torch.Tensor, + conv_state: torch.Tensor, + starts: torch.Tensor, + slots: torch.Tensor, + has_init: torch.Tensor, + max_seqlen: int, + ) -> None: + # The op checks dtypes, shapes, contiguity and storage overlap itself; it does not read the metadata + # values on the host, and these are the ones it gets wrong silently. + bounds = starts.tolist() + assert ( + bounds[0] == 0 + and bounds[-1] == src.shape[0] + and all(lo <= hi for lo, hi in zip(bounds, bounds[1:])) + ), ( + "starts must monotonically partition the packed tokens: uncovered rows of q / k / v stay unwritten" + ) + slot_list = slots.tolist() + assert len(set(slot_list)) == len(slot_list), ( + "slots must be unique: two sequences on one slot race on its history" + ) + assert all(0 <= slot < conv_state.shape[0] for slot in slot_list), ( + "slots must lie in the pool: slot -1 copies the raw inputs through unconvolved and other " + "negative slots index the pool from its end" + ) + + +kda_prefill_conv = KdaPrefillConv() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/ssm/reset_recurrent_state_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/ssm/reset_recurrent_state_rows.py new file mode 100644 index 000000000000..a5f5c87dea62 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ops/ssm/reset_recurrent_state_rows.py @@ -0,0 +1,153 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Zero the recurrent-state rows of requests that start fresh. + +One program per (request, 256-element block of the row), masked by the request's flag and slot. Modeling +V1's ``reset_recurrent_state_rows`` (``tensorrt_llm/_torch/modules/mamba/recurrent_state_cache.py``) +launches the same kernel, which can also zero the conv-state rows; the target's fused prefill convolution +clears fresh conv rows itself, so that path is not part of this entry. This entry does not import V1. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Elements of a row one program zeroes. +_BLOCK_SIZE = 256 + + +@triton.jit +def _reset_recurrent_state_rows_kernel( + recurrent_states, + state_indices, + has_initial_states, + recurrent_state_stride, + NUM_CACHE_LINES: tl.constexpr, + RECURRENT_STATE_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + request_idx = tl.program_id(0) + offsets = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + state_idx = tl.load(state_indices + request_idx).to(tl.int64) + needs_reset = ~tl.load(has_initial_states + request_idx).to(tl.int1) + valid_state = (state_idx >= 0) & (state_idx < NUM_CACHE_LINES) + recurrent_row_offset = state_idx * recurrent_state_stride.to(tl.int64) + tl.store( + recurrent_states + recurrent_row_offset + offsets, + 0.0, + mask=needs_reset & valid_state & (offsets < RECURRENT_STATE_SIZE), + ) + + +class ResetRecurrentStateRows(OpWrapper): + """Zero slot ``state_indices[i]``'s row of ``recurrent_states`` for every request i with + ``has_initial_states[i]`` false; one launch, returns None. + + Args: + recurrent_states: ``[slots, ...]`` pool (KDA: ``[slots, H, V, K]`` bf16 or fp32), each slot's row + contiguous; the slot stride is free (the cache manager's per-slot records are wider than one row). + state_indices: int32 or int64 ``[N]`` slot of each request. Indices outside ``[0, slots)`` are skipped. + has_initial_states: bool ``[N]``; a request with True keeps its row. + """ + + ARCHS = frozenset({Arch.SM_100}) + # GLM-5.3-Flash's KDA at TP4 (16 of 64 heads, K = V = 128), on pools whose slot stride is a whole per-slot + # record, as the hybrid cache manager lays them out. The shapes the shipped target resets are in this entry's + # test. + CELLS: tuple[Cell, ...] = ( + Cell( + why="fp32 pool: out-of-range and negative slots skipped", + spec=dict( + slots=10, + indices=(3, -1, 9, 10, 0), + has_init=(False, False, True, False, False), + state_dtype=torch.float32, + ), + ), + Cell( + why="int64 slot indices, every request fresh, rows larger than one 256-element program", + spec=dict( + slots=6, + indices=(5, 1, 2), + has_init=(False, False, False), + state_dtype=torch.bfloat16, + index_dtype=torch.int64, + ), + ), + ) + note = """ + The flag is read per request: True keeps the slot's row, False zeroes it. A + slot index outside [0, slots) (the -1 of a padding request, or a slot past the + pool) is skipped rather than rejected. Zeroing is a masked store of 0.0, so any + pool dtype works; only the elements of the slot's row are written, never the + padding between rows. + + KDA prefill calls this before the recurrence: the prefill kernel reads every + selected slot as its initial state, so a fresh request's slot must be zeroed + first. + + The launch is ordered on the current stream and takes no host sync, so it is + CUDA-graph capturable once compiled; the row size and slot count are + compile-time constants, so the first call per pool geometry compiles. + """ + + def __call__( + self, + recurrent_states: torch.Tensor, + state_indices: torch.Tensor, + has_initial_states: torch.Tensor, + ) -> None: + num_requests = state_indices.shape[0] + if num_requests == 0: + return + recurrent_state_size = recurrent_states.numel() // recurrent_states.shape[0] + grid = (num_requests, triton.cdiv(recurrent_state_size, _BLOCK_SIZE)) + _reset_recurrent_state_rows_kernel[grid]( + recurrent_states, + state_indices, + has_initial_states, + recurrent_states.stride(0), + recurrent_states.shape[0], + recurrent_state_size, + _BLOCK_SIZE, + ) + + def reference( + self, + recurrent_states: torch.Tensor, + state_indices: torch.Tensor, + has_initial_states: torch.Tensor, + ) -> torch.Tensor: + """``recurrent_states`` as the call leaves it; mutates nothing. Exact: the op writes zeros or nothing.""" + recurrent = recurrent_states.clone() + for slot, keep in zip(state_indices.tolist(), has_initial_states.tolist()): + if keep or not 0 <= slot < recurrent.shape[0]: + continue + recurrent[slot] = 0 + return recurrent + + def is_valid( + self, + recurrent_states: torch.Tensor, + state_indices: torch.Tensor, + has_initial_states: torch.Tensor, + ) -> None: + # The kernel zeroes numel / slots elements from slot * stride(0) and checks nothing. + assert recurrent_states.shape[0] > 0, "the recurrent pool must have at least one slot" + assert recurrent_states[0].is_contiguous(), ( + "each recurrent-state row must be contiguous: the kernel zeroes one dense run per slot" + ) + assert recurrent_states.stride(0) >= recurrent_states[0].numel(), ( + "recurrent-state rows must not overlap: zeroing one would clear part of the next" + ) + assert state_indices.ndim == 1, "state_indices must be one slot per request" + assert has_initial_states.shape == state_indices.shape, ( + "has_initial_states must hold one flag per request: the kernel reads it by request index" + ) + + +reset_recurrent_state_rows = ResetRecurrentStateRows() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/schedule.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/schedule.py new file mode 100644 index 000000000000..32d9aa6af247 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/schedule.py @@ -0,0 +1,219 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""What one forward step looks like, derived once and shared by every layer of the step. + +The engine prepares the attention metadata (``TrtllmAttentionMetadata``) and, inside it, +the hybrid-cache metadata of this model (``mamba_metadata``, the class the shell declares +as ``mamba_metadata_cls``): device lengths, the GLM block tables and the host schedules of +the context requests. Everything here reads those prepared fields; nothing reconstructs +cache tables or lengths, so a step captured in a CUDA graph replays on the refreshed +buffers. + +A step's rows are packed context first: ``num_ctx_tokens`` rows of ``num_contexts`` +context requests, then the generation rows, ``gen_tokens_per_request`` per generation +request (1 for plain decode, 1 + the draft length while MTP verifies drafts). +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Optional, Sequence + +import torch +from transformers import PretrainedConfig + +from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata +from tensorrt_llm._utils import async_tensor_h2d + +LINEAR_ATTENTION = "linear_attention" +SPARSE_ATTENTION = "deepseek_sparse_attention" +DENSE_MLP = "dense" +SPARSE_MLP = "sparse" + + +@dataclass(frozen=True) +class LayerSchedule: + """The checkpoint's per-layer attention and MLP kinds (``layer_types``, + ``mlp_layer_types``), cross-checked against the redundant config fields.""" + + attention: tuple[str, ...] + mlp: tuple[str, ...] + + @property + def num_layers(self) -> int: + return len(self.attention) + + @classmethod + def from_text_config(cls, text: PretrainedConfig) -> "LayerSchedule": + num_layers = int(text.num_hidden_layers) + attention = tuple(text.layer_types) + mlp = tuple(text.mlp_layer_types) + for name, values, allowed in ( + ("layer_types", attention, {LINEAR_ATTENTION, SPARSE_ATTENTION}), + ("mlp_layer_types", mlp, {DENSE_MLP, SPARSE_MLP}), + ): + if len(values) != num_layers or set(values) - allowed: + raise ValueError(f"{name} {values} does not list {num_layers} known layer kinds") + schedule = cls(attention, mlp) + linear = dict(getattr(text, "linear_attn_config", None) or {}) + kda = tuple(i for i, t in enumerate(attention) if t == LINEAR_ATTENTION) + if linear.get("kda_layers") is not None and tuple(linear["kda_layers"]) != kda: + raise ValueError("linear_attn_config.kda_layers disagrees with layer_types") + first_dense = getattr(text, "first_k_dense_replace", None) + dense = tuple(i for i, t in enumerate(mlp) if t == DENSE_MLP) + if first_dense is not None and dense != tuple(range(int(first_dense))): + raise ValueError("mlp_layer_types disagrees with first_k_dense_replace") + return schedule + + +@dataclass(frozen=True) +class ContextRows: + """Device row schedule of the packed context tokens of one step. + + ``positions[i]`` is packed context token ``i``'s cache position and ``request_ids[i]`` + its request (the batch row of its block table). ``final_positions`` / + ``final_request_ids`` list the pool-final positions this step completes, one per pool, + so pool-key refreshes never race. + """ + + positions: torch.Tensor # [num_ctx_tokens] int64 + request_ids: torch.Tensor # [num_ctx_tokens] int32 + final_positions: torch.Tensor # [F] int64 + final_request_ids: torch.Tensor # [F] int32 + + @classmethod + def build( + cls, cu_seqlens: Sequence[int], cached_lens: Sequence[int], kpool: int, device: torch.device + ) -> "ContextRows": + positions: list[int] = [] + request_ids: list[int] = [] + finals: list[int] = [] + final_ids: list[int] = [] + for i in range(len(cu_seqlens) - 1): + length = int(cu_seqlens[i + 1]) - int(cu_seqlens[i]) + cached = int(cached_lens[i]) + positions.extend(range(cached, cached + length)) + request_ids.extend([i] * length) + first_final = cached + (-(cached + 1)) % kpool + for pos in range(first_final, cached + length, kpool): + finals.append(pos) + final_ids.append(i) + # Pinned non-blocking copies do not wait for the GPU work already queued. + return cls( + positions=async_tensor_h2d(positions, torch.long, device), + request_ids=async_tensor_h2d(request_ids, torch.int32, device), + final_positions=async_tensor_h2d(finals, torch.long, device), + final_request_ids=async_tensor_h2d(final_ids, torch.int32, device), + ) + + +@dataclass(frozen=True) +class GenerationRows: + """Device row schedule of the generation requests of one step, built from the live + ``kv_lens`` only (so CUDA graphs replay it on refreshed metadata). + + Row ``b * T + t`` is generation request ``b``'s token ``t``: ``positions`` is ``[B, T]``, + ``visible`` (the row's visible prefix length) is ``[B * T]``, and ``tables`` is the + request's block table repeated per row, as the pool scoring and expansion kernels address + it. + """ + + positions: torch.Tensor # [B, T] int64 + visible: torch.Tensor # [B * T] int64 + tables: torch.Tensor # [B * T, max pages] int64 + + @classmethod + def build( + cls, kv_lens: torch.Tensor, block_tables: torch.Tensor, tokens: int + ) -> "GenerationRows": + steps = torch.arange(tokens, device=kv_lens.device) + positions = (kv_lens - tokens).unsqueeze(1) + steps + visible = (positions + 1).reshape(-1) + return cls(positions, visible, block_tables.repeat_interleave(tokens, dim=0)) + + +@dataclass +class StepSchedule: + """Per-step facts and row schedules shared by every layer of one forward.""" + + num_contexts: int + num_ctx_tokens: int + num_generations: int + #: Tokens per generation request: 1 for plain decode, 1 + draft length while verifying. + gen_tokens_per_request: int + #: Host packed offsets and cached lengths of the context requests (prefill only). + ctx_cu_seqlens: list[int] + cached_lens: list[int] + #: Live per-request visible lengths (cached + this step's tokens) on the device. + kv_lens: torch.Tensor + #: The GLM block tables of the batch, ``[batch, max pages]`` int64 slot ids. + block_tables: torch.Tensor + metadata: AttentionMetadata + _context_rows: dict[int, ContextRows] = field(default_factory=dict) + _generation_rows: Optional[GenerationRows] = None + + @property + def num_generation_tokens(self) -> int: + return self.num_generations * self.gen_tokens_per_request + + def context_rows(self, kpool: int, device: torch.device) -> ContextRows: + rows = self._context_rows.get(kpool) + if rows is None: + rows = ContextRows.build( + self.ctx_cu_seqlens, self.cached_lens[: self.num_contexts], kpool, device + ) + self._context_rows[kpool] = rows + return rows + + def generation_rows(self) -> GenerationRows: + if self._generation_rows is None: + self._generation_rows = GenerationRows.build( + self.kv_lens[self.num_contexts :], + self.block_tables[self.num_contexts : self.num_contexts + self.num_generations], + self.gen_tokens_per_request, + ) + return self._generation_rows + + @classmethod + def build(cls, attn_metadata: AttentionMetadata) -> "StepSchedule": + """The step's schedule from the engine-prepared metadata. + + Prefill reads the host context schedules the metadata prepared before the forward; + decode and verify read ``kv_lens_cuda``, which carries the overlap scheduler's + corrections and MTP's rewinds on the device. + """ + if attn_metadata.kv_cache_manager is None: + raise ValueError("GLM-5.3-Flash needs a KV cache manager") + hybrid = attn_metadata.mamba_metadata + if hybrid is None or getattr(hybrid, "glm_block_tables", None) is None: + raise RuntimeError( + "GLM-5.3-Flash needs the prepared GLM hybrid-cache metadata (glm_block_tables); " + "prepare the attention metadata before eager execution or CUDA graph capture" + ) + live = attn_metadata.kv_lens_cuda + if live is None: + raise ValueError("GLM-5.3-Flash needs the prepared kv_lens_cuda") + batch = int(attn_metadata.seq_lens.shape[0]) + num_contexts = int(attn_metadata.num_contexts) + num_generations = batch - num_contexts + num_ctx_tokens = int(attn_metadata.num_ctx_tokens) + tokens = 1 + if num_generations > 0: + gen_tokens = int(attn_metadata.num_tokens) - num_ctx_tokens + if gen_tokens <= 0 or gen_tokens % num_generations: + raise ValueError( + f"{gen_tokens} generation tokens do not split evenly over " + f"{num_generations} generation requests" + ) + tokens = gen_tokens // num_generations + return cls( + num_contexts=num_contexts, + num_ctx_tokens=num_ctx_tokens, + num_generations=num_generations, + gen_tokens_per_request=tokens, + ctx_cu_seqlens=hybrid.glm_ctx_cu_seqlens, + cached_lens=hybrid.glm_cached_lens_host, + kv_lens=live[:batch], + block_tables=hybrid.glm_block_tables[:batch], + metadata=attn_metadata, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/weights.py new file mode 100644 index 000000000000..ebecebc5610b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/weights.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""This rank's weights, read from the checkpoint the engine hands ``load_weights``. + +Every tensor is read at this rank's share (``checkpoint.Checkpoint.read`` cuts the lazy +safetensors entry, so no rank reads another's rows), transformed into the layout the +target's operations read, and kept as a plain device tensor on the core. Under tensor +parallelism 4: + +* **KDA** projections, convolution taps, gates and the output projection are split by this + rank's 16 of 64 heads (``layers/kda.py``). +* **Sparse MLA**: the query heads (``q_b_proj``, ``kv_b_proj``) and the matching input + columns of ``o_proj`` are split by 16 of 64 heads; the ``[q_a | kv_a]`` latent + projections, the norms and the whole indexer are replicated, as are the latent and + indexer cache rows they write. Sequence-parallel prefill attends a quarter of the rows + with all 64 heads, so the layer also keeps full-head copies (``layers/mla.py``). +* **FFN**: the dense MLPs' and the experts' intermediate is split (the routed experts by + moe_tp 4, all 288 on every rank); router and bias are replicated (``layers/ffn.py``). +* Everything here is replicated: the embedding, every norm, the mHC mixers, and the MTP + layer's ``enorm`` / ``hnorm`` / ``shared_head.norm``. ``eh_proj`` is row-parallel (this + rank's 2048 of its 8192 input columns). ``lm_head`` is the shell's vocabulary-parallel + head: this rank's rows of the checkpoint tensor. + +The checkpoint is block FP8 (128 x 128 blocks, FP32 ``weight_scale_inv``) for the MLA +low-rank and query / output projections and every MLP and expert projection, of the MTP +layer too; the KDA layers, ``kv_b_proj``, the indexer, the router, the norms, ``eh_proj`` +and the mHC mixers are BF16 or FP32. ``activation_scheme`` is dynamic: there are no input +scales. + +Loading contract: every checkpoint key is read exactly once, except the vision tower's, +and the MTP layer's when no speculative decoding builds it. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Mapping, Optional, Union + +import torch + +from .checkpoint import Checkpoint, cols, rows, share +from .dims import TEXT_PREFIX, Dims +from .layers.ffn import DenseMlpWeights, MoeWeights, load_dense_mlp, load_moe +from .layers.kda import KdaWeights, load_kda +from .layers.mla import MlaWeights, load_mla +from .schedule import DENSE_MLP, LINEAR_ATTENTION + +#: Checkpoint key prefix of the vision tower, which this text-only target never loads. +VISION_PREFIX = "model.visual." + + +@dataclass +class HcWeights: + """One mHC site's mixer: ``fn`` FP32 ``[mix, mult * hidden]`` (the checkpoint's BF16 cast + to FP32), ``base`` FP32 ``[mix]`` and ``scale`` FP32 ``[3]``.""" + + fn: torch.Tensor + base: torch.Tensor + scale: torch.Tensor + + +@dataclass +class DecoderLayerWeights: + attention_kind: str + mlp_kind: str + input_norm: torch.Tensor # [hidden] bf16 + post_norm: torch.Tensor # [hidden] bf16 + hc_attn: HcWeights + hc_ffn: HcWeights + attention: Union[KdaWeights, MlaWeights] + mlp: Union[DenseMlpWeights, MoeWeights] + + +@dataclass +class MtpLayerWeights: + """The MTP layer (checkpoint layer ``num_hidden_layers``): a sparse-MLA + MoE layer + without mHC, plus its front end and head norm.""" + + enorm: torch.Tensor # [hidden] bf16 + hnorm: torch.Tensor # [hidden] bf16 + eh_proj: torch.Tensor # [hidden, 2 * hidden / tp] bf16, this rank's input columns + head_norm: torch.Tensor # [hidden] bf16, shared_head.norm + input_norm: torch.Tensor # [hidden] bf16 + post_norm: torch.Tensor # [hidden] bf16 + attention: MlaWeights + mlp: MoeWeights + + +@dataclass +class ModelWeights: + embed: torch.Tensor # [vocab, hidden] bf16, replicated + final_norm: torch.Tensor # [hidden] bf16 + layers: list[DecoderLayerWeights] + mtp: Optional[MtpLayerWeights] + + +def _hc(ck: Checkpoint, prefix: str, site: str) -> HcWeights: + return HcWeights( + fn=ck.read(f"{prefix}hc_{site}_fn", dtype=torch.float32), + base=ck.read(f"{prefix}hc_{site}_base", dtype=torch.float32), + scale=ck.read(f"{prefix}hc_{site}_scale", dtype=torch.float32), + ) + + +def _decoder_layer(ck: Checkpoint, dims: Dims, layer: int) -> DecoderLayerWeights: + prefix = dims.layer_prefix(layer) + attention_kind = dims.schedule.attention[layer] + mlp_kind = dims.schedule.mlp[layer] + attention = ( + load_kda(ck, dims, layer) + if attention_kind == LINEAR_ATTENTION + else load_mla(ck, dims, layer) + ) + mlp = load_dense_mlp(ck, dims, layer) if mlp_kind == DENSE_MLP else load_moe(ck, dims, layer) + return DecoderLayerWeights( + attention_kind=attention_kind, + mlp_kind=mlp_kind, + input_norm=ck.read(f"{prefix}input_layernorm.weight"), + post_norm=ck.read(f"{prefix}post_attention_layernorm.weight"), + hc_attn=_hc(ck, prefix, "attn"), + hc_ffn=_hc(ck, prefix, "ffn"), + attention=attention, + mlp=mlp, + ) + + +def _mtp_layer(ck: Checkpoint, dims: Dims) -> MtpLayerWeights: + layer = dims.mtp_layer + prefix = dims.layer_prefix(layer) + start, end = share(2 * dims.hidden, dims.tp_size, dims.tp_rank) + return MtpLayerWeights( + enorm=ck.read(f"{prefix}enorm.weight"), + hnorm=ck.read(f"{prefix}hnorm.weight"), + eh_proj=ck.read(f"{prefix}eh_proj.weight", cut=cols(start, end)), + head_norm=ck.read(f"{prefix}shared_head.norm.weight"), + input_norm=ck.read(f"{prefix}input_layernorm.weight"), + post_norm=ck.read(f"{prefix}post_attention_layernorm.weight"), + attention=load_mla(ck, dims, layer), + mlp=load_moe(ck, dims, layer), + ) + + +def load_model_weights( + weights: Mapping[str, Any], dims: Dims, lm_head_weight: torch.Tensor, device: torch.device +) -> ModelWeights: + """Read this rank's weights and fill the shell's vocabulary-parallel ``lm_head`` share.""" + ck = Checkpoint(weights, device) + model = ModelWeights( + embed=ck.read(f"{TEXT_PREFIX}embed_tokens.weight"), + final_norm=ck.read(f"{TEXT_PREFIX}norm.weight"), + layers=[_decoder_layer(ck, dims, i) for i in range(dims.num_layers)], + mtp=_mtp_layer(ck, dims) if dims.mtp else None, + ) + start = dims.tp_rank * dims.vocab_local + end = min(start + dims.vocab_local, dims.vocab) + head = ck.read("lm_head.weight", cut=rows(start, end)) + assert lm_head_weight.shape[1] == dims.hidden and lm_head_weight.shape[0] >= end - start, ( + f"lm_head share {tuple(lm_head_weight.shape)} for rows [{start}, {end})" + ) + with torch.no_grad(): + lm_head_weight[: end - start].copy_(head) + lm_head_weight[end - start :].zero_() + skipped = {k for k in weights.keys() if k.startswith(VISION_PREFIX)} + if not dims.mtp: + mtp_prefix = dims.layer_prefix(dims.mtp_layer) + skipped |= {k for k in weights.keys() if k.startswith(mtp_prefix)} + ck.check_coverage(skipped) + return model diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/routing.py index 168dbea9a447..7382c3d1f3d3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/routing.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next/routing.py @@ -8,7 +8,7 @@ uses the built-in GLM-5.3-Flash implementation; in ``require`` it raises that error. -Two targets share one checkpoint, GPU and parallel coverage and differ in how +Two NVFP4 targets share one checkpoint, GPU and parallel coverage and differ in how speculative decoding runs: * ``glm_5_3_flash_nvfp4__sm_100__default_mtp_fused`` runs one-model MTP-Eagle @@ -17,6 +17,11 @@ speculative worker runs after its forward. It takes every other deployment it supports, including no speculative decoding. +The checkpoint's quantization picks the targets that can apply: the two above +serve the NVFP4 checkpoint, and ``glm_5_3_flash_fp8__sm_10x__tp4`` serves the +block-FP8 (``FP8_BLOCK_SCALES``) checkpoint on sm_100 with tensor parallelism 4 +(see ``_fp8_refused``). + The checkpoint declares the vision-language wrapper, and the language model's shape lives in its ``text_config``, so that is what this tree reads. The targets are text-only: they load no vision tower. @@ -24,19 +29,25 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Optional, Tuple +from typing import TYPE_CHECKING, Any, Iterable, Optional, Tuple -from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig, SADecodingConfig +from tensorrt_llm.llmapi.llm_args import ( + DecodingBaseConfig, + LoadFormat, + MTPDecodingConfig, + PrefillCudaGraphBackend, + SADecodingConfig, +) from tensorrt_llm.quantization.mode import QuantAlgo -from ..._router_index import NoTarget, routes +from ..._router_index import NoTarget, engine_kv_cache_quant_algo, routes if TYPE_CHECKING: from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm.llmapi.llm_args import TorchLlmArgs -# The one GPU architecture the targets are written for: B200 (sm100). They have -# been run on nothing else, so routing sends nothing else to them. +# The one GPU architecture the targets are written for: B200 (sm100). The NVFP4 +# targets have been run on nothing else, so routing sends nothing else to them. _SM = (10, 0) # GLM-5.3-Flash's config shape, read from ``text_config``: (num_hidden_layers, @@ -51,6 +62,34 @@ "linear_num_heads", ) +# The GPU architecture the block-FP8 target has run on: sm_100 (B200, GB200). It reads +# the same config shape as the NVFP4 targets, and the geometry below. +_FP8_SM = (10, 0) + +# What the block-FP8 target's kernels compile in beyond that shape: the experts each token +# routes to, the experts' intermediate size (512 per rank in the decode MoE kernel) and the +# number of mHC streams. +_FP8_GEOMETRY = {"num_experts_per_tok": 8, "moe_intermediate_size": 2048, "hc_mult": 4} + +# The KDA head size and conv width the block-FP8 target's KDA kernels implement. +_FP8_KDA = {"head_dim": 128, "short_conv_kernel_size": 4} + +#: The MoE routing the block-FP8 target's decode, prefill and runner kernels implement +#: (DeepSeek-V3's noaux_tc): sigmoid scores, top-k over the scores plus the correction bias +#: in a single group, the selected scores normalized, and one shared expert. +FP8_MOE_ROUTING = { + "scoring_func": "sigmoid", + "topk_method": "noaux_tc", + "n_group": 1, + "topk_group": 1, + "norm_topk_prob": True, + "n_shared_experts": 1, +} + +# The all-reduce strategies the block-FP8 target runs: the one-shot all-reduce, whose +# workspace its fused decode boundary shares, and NCCL above the one-shot row limit. +_FP8_ALLREDUCE = ("AUTO", "ONESHOT") + def _field(config: Any, name: str) -> Any: """Read ``name`` off a sub-config that may be an object or a plain dict.""" @@ -159,7 +198,7 @@ def route( None if text_config is None else tuple(_field(text_config, name) for name in _SHAPE_FIELDS) ) quant_algo = _quant_algo(model_config) - # Both targets compute their weight shapes from the mapping, as the built-in + # The NVFP4 targets compute their weight shapes from the mapping, as the built-in # model does; neither runs pipeline or context parallelism. single_stage = mapping.pp_size == 1 and mapping.cp_size == 1 deployment = ( @@ -167,10 +206,20 @@ def route( f"quantization {quant_algo.name if quant_algo is not None else 'none'}, " f"pp={mapping.pp_size} cp={mapping.cp_size}" ) + if quant_algo == QuantAlgo.FP8_BLOCK_SCALES: + fp8_refused = _fp8_refused(model_config, llm_args, compute_capability) + if fp8_refused: + raise NoTarget( + f"{deployment}; the block-FP8 glm5_next target does not run " + f"{'; '.join(fp8_refused)}" + ) + from .glm_5_3_flash_fp8__sm_10x__tp4.modeling import ModelingV2Glm53FlashFp8Sm10xTp4 + + return ModelingV2Glm53FlashFp8Sm10xTp4 if not (compute_capability == _SM and shape == _SHAPE and quant_algo == QuantAlgo.NVFP4): raise NoTarget( f"{deployment}; the glm5_next targets are the NVFP4 GLM-5.3-Flash checkpoint on " - "sm100 (B200)" + "sm100 (B200) and the block-FP8 (FP8_BLOCK_SCALES) one on sm100 / sm103" ) if not single_stage: raise NoTarget( @@ -199,3 +248,142 @@ def route( "targets run one-model MTP (vanilla or MTP-Eagle, from the checkpoint's MTP layer), " "suffix-automaton drafting or no speculative decoding" ) + + +def _fp8_speculation_refused(spec_config: Optional[DecodingBaseConfig]) -> list[str]: + """The speculative features the block-FP8 target does not run; empty when it runs this config. + + It drafts with the checkpoint's own MTP layer, driven by the engine's one-model MTP-Eagle + worker at three linear drafts with strict acceptance. Its KDA layers verify exactly 1 + 3 + tokens per request, so a draft length that changes between steps is declined, and its + draft embedding reads token ids only, never the -1 padding rejection sampling leaves in + the accepted tokens. A separate draft checkpoint, a draft tree, an MXFP8 draft head and a + draft MoE backend, none of which it builds, are declined rather than ignored. + """ + if spec_config is None: + return [] + if not isinstance(spec_config, MTPDecodingConfig): + return [f"{type(spec_config).__name__} (only MTP)"] + refused = { + f"{spec_config.spec_dec_mode.name} (only one-model MTP-Eagle)": ( + not spec_config.spec_dec_mode.is_mtp_eagle_one_model() + ), + f"max_draft_len={spec_config.max_draft_len} (only 3)": spec_config.max_draft_len != 3, + "draft trees": spec_config.max_total_draft_tokens != spec_config.max_draft_len, + "a dynamic draft tree": spec_config.use_dynamic_tree, + # max_concurrency is translated into a draft-length schedule. + "a draft-length schedule": spec_config.draft_len_schedule is not None, + # The executor builds this gate under the same condition; once it trips, every later + # step runs with no drafts. + "the acceptance-rate gate": bool(spec_config.acceptance_rate_window_size) + and spec_config.acceptance_rate_threshold is not None, + "rejection sampling": spec_config.use_rejection_sampling, + "relaxed acceptance": spec_config.use_relaxed_acceptance_for_thinking, + "a separate speculative checkpoint": spec_config.speculative_model is not None, + f"draft_lm_head_dtype={spec_config.draft_lm_head_dtype} (only auto)": ( + spec_config.draft_lm_head_dtype != "auto" + ), + f"draft MoE backend {spec_config.moe_backend}": spec_config.moe_backend is not None, + } + return [name for name, present in refused.items() if present] + + +def fp8_model_refused(model_config: ModelConfig, compute_capability: Tuple[int, int]) -> list[str]: + """What the block-FP8 target does not run among the facts ``model_config`` states; empty when + it runs them. + + The GPU, the checkpoint's shape and geometry, the mapping and the speculative configuration. + ``route`` declines these together with the engine settings in :func:`_fp8_refused`, and the + target's constructor raises on them for a direct construction. + """ + text_config = getattr(model_config.pretrained_config, "text_config", None) + mapping = model_config.mapping + + def fields(names: Iterable[str]) -> Optional[dict[str, Any]]: + return None if text_config is None else {name: _field(text_config, name) for name in names} + + shape = None if text_config is None else tuple(_field(text_config, n) for n in _SHAPE_FIELDS) + geometry = fields(_FP8_GEOMETRY) + routing = fields(FP8_MOE_ROUTING) + linear = None if text_config is None else _field(text_config, "linear_attn_config") + gate_lower_bound = None if linear is None else _field(linear, "gate_lower_bound") + kda = None if linear is None else {name: _field(linear, name) for name in _FP8_KDA} + tied = text_config is not None and bool(_field(text_config, "tie_word_embeddings")) + refused = { + f"compute capability {compute_capability} (only sm_100)": compute_capability != _FP8_SM, + f"text_config shape {shape} (only {_SHAPE})": shape != _SHAPE, + f"text_config geometry {geometry} (only {_FP8_GEOMETRY})": geometry != _FP8_GEOMETRY, + f"MoE routing {routing} (only {FP8_MOE_ROUTING})": routing != FP8_MOE_ROUTING, + "a KDA gate without gate_lower_bound": gate_lower_bound is None, + f"KDA {kda} (only {_FP8_KDA})": kda != _FP8_KDA, + "an LM head tied to the embedding": tied, + f"world_size={mapping.world_size} tp={mapping.tp_size} (only tp4 over 4 ranks)": not ( + mapping.world_size == 4 and mapping.tp_size == 4 + ), + f"moe_tp={mapping.moe_tp_size} moe_ep={mapping.moe_ep_size} (only the experts split by " + "intermediate, moe_tp=4 moe_ep=1)": not ( + mapping.moe_tp_size == 4 and mapping.moe_ep_size == 1 + ), + "attention DP": mapping.enable_attention_dp, + f"pp={mapping.pp_size} cp={mapping.cp_size} (only one stage)": ( + mapping.pp_size != 1 or mapping.cp_size != 1 + ), + f"gpus_per_node={mapping.gpus_per_node} (only the 4 ranks of one node)": ( + mapping.gpus_per_node < 4 + ), + } + return [name for name, present in refused.items() if present] + _fp8_speculation_refused( + model_config.spec_config + ) + + +def _fp8_refused( + model_config: ModelConfig, + llm_args: TorchLlmArgs, + compute_capability: Tuple[int, int], +) -> list[str]: + """What the block-FP8 target does not run for this deployment; empty when it is admitted. + + The target serves GLM-5.3-Flash's block-FP8 release, text only, on sm_100 with tp4: + tensor-parallel attention over the 4 ranks of one node and the routed experts split by + intermediate, a weight layout, which is why it selects a target rather than a runtime branch + (:func:`fp8_model_refused`). Its sparse attention reads an FP8 latent cache. Its MoE kernels + read the routed experts in the checkpoint's layout, which only the TRTLLM backend keeps (AUTO + resolves to it here), run every row of a step in one call, and do no load balancing or DWDP. + Its forward builds the step's schedule on the host, so it runs no prefill CUDA graphs and is + not compiled. It loads the checkpoint's weights, computes the logits in BF16 and runs its + collectives on the one-shot all-reduce (NCCL above the one-shot row limit). Each of these is + declined rather than ignored. + """ + moe = llm_args.moe_config + kv_cache_quant_algo = engine_kv_cache_quant_algo(model_config, llm_args) + kv_cache = ( + "bf16" + if kv_cache_quant_algo is None + else getattr(kv_cache_quant_algo, "name", kv_cache_quant_algo) + ) + load_format = getattr(llm_args.load_format, "name", llm_args.load_format) + refused = { + f"KV cache {kv_cache} (only FP8)": kv_cache_quant_algo != QuantAlgo.FP8, + "the vision encoder (only disable_mm_encoder)": not llm_args.disable_mm_encoder, + f"MoE backend {moe.backend} (only AUTO or TRTLLM)": moe.backend not in ("AUTO", "TRTLLM"), + f"a MoE load balancer ({moe.load_balancer})": moe.load_balancer is not None, + f"moe_config.max_num_tokens={moe.max_num_tokens} (the MoE runs every row of a step)": ( + moe.max_num_tokens is not None + ), + f"DWDP ({llm_args.dwdp_config})": llm_args.dwdp_config is not None, + f"prefill CUDA graphs ({llm_args.prefill_cuda_graph_backend})": ( + llm_args.prefill_cuda_graph_backend != PrefillCudaGraphBackend.DISABLED + ), + "torch.compile": llm_args.torch_compile_config is not None, + f"load_format {load_format} (only AUTO)": llm_args.load_format != LoadFormat.AUTO, + f"lm_head_output_dtype={llm_args.lm_head_output_dtype} (only auto)": ( + llm_args.lm_head_output_dtype != "auto" + ), + f"allreduce_strategy={llm_args.allreduce_strategy} (only AUTO or ONESHOT)": ( + llm_args.allreduce_strategy not in _FP8_ALLREDUCE + ), + } + return fp8_model_refused(model_config, compute_capability) + [ + name for name, present in refused.items() if present + ] diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py index 1e4d7da8c654..9ab0e2660af4 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py @@ -453,9 +453,11 @@ def promote_final_contexts( and not ced_enabled and self._config.max_beam_width == 1 and self._context_promotion_supported - # PLE owns recurrent n-gram and convolution state. Promoting a - # fresh final-context row would skip its cache-slot reset. + # PLE owns recurrent n-gram and convolution state, and other models + # may declare state that only their context path resets. Promoting + # a fresh final-context row would skip its cache-slot reset. and not self._model_uses_ple_recurrent_state + and not self._model_resets_recurrent_state_in_context and self.mapping.cp_size == 1 ): return scheduled_requests, frozenset() @@ -483,6 +485,21 @@ def _is_final_multimodal_context_decode_compatible(self, request: LlmRequest) -> return True return CUDAGraphRunner._get_mrope_position_delta(request) is not None + @functools.cached_property + def _model_resets_recurrent_state_in_context(self) -> bool: + """Whether the model declares recurrent state that only its context path resets. + + A model sets ``resets_recurrent_state_in_context = True`` when a fresh + cache slot is reset by its context forward and not by its decode forward. + The declaration may sit on the causal-LM wrapper the engine holds or on + the decoder inside it. + """ + wrapper = getattr(self.model, "_orig_mod", self.model) + return any( + bool(getattr(model, "resets_recurrent_state_in_context", False)) + for model in (wrapper, get_top_level_model(self.model)) + ) + @functools.cached_property def _model_uses_ple_recurrent_state(self) -> bool: """Detect PLE on text-only and multimodal model wrappers. diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index 09fda5f35066..6d959f3f40c1 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -536,6 +536,7 @@ def apply_model_defaults_to_llm_args( "via get_preferred_transceiver_runtime() instead.") user_overrides = llm_args.model_dump(exclude_unset=True) + fields_set = set(llm_args.model_fields_set) base_state = llm_args.model_dump() merged_state = _deep_merge(base_state, model_defaults_dict, user_overrides) @@ -563,7 +564,13 @@ def _compute_applied(defaults: Dict[str, Any], applied[key] = default_value return applied - return _compute_applied(model_defaults_dict, user_overrides) + applied = _compute_applied(model_defaults_dict, user_overrides) + # Assigning every field above marks each one as set. Keep only the fields + # the user set and the ones the model defaults applied, so checks of + # whether a deployment set a field (``model_fields_set``) still hold. + object.__setattr__(llm_args, "__pydantic_fields_set__", + fields_set | applied.keys()) + return applied def _resolve_kv_cache_manager_v2_auto(llm_args: 'TorchLlmArgs', diff --git a/tests/integration/defs/accuracy/modeling_v2/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/__init__.py b/tests/integration/defs/accuracy/modeling_v2/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tests/integration/defs/accuracy/modeling_v2/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/__init__.py @@ -0,0 +1,14 @@ +# 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. diff --git a/tests/integration/defs/accuracy/modeling_v2/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/test_accuracy.py b/tests/integration/defs/accuracy/modeling_v2/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/test_accuracy.py new file mode 100644 index 000000000000..cdd0b5ab9641 --- /dev/null +++ b/tests/integration/defs/accuracy/modeling_v2/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/test_accuracy.py @@ -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. +"""Whole-model gates for the glm5_next modeling_v2 target glm_5_3_flash_fp8__sm_10x__tp4. + +Every test passes ``modeling_v2="require"``: under ``"auto"`` a deployment the target does +not claim would fall back to the built-in model, pass, and report the built-in model's +numbers as the target's. The settings are the target's routed deployment: tensor +parallelism 4 with the experts split by intermediate, an FP8 latent cache, the vision tower +disabled. +""" + +import pytest + +from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm._torch._experimental.modeling_v2 import MODELING_V2_VALIDATE_ENV +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, MTPDecodingConfig + +from .....conftest import llm_models_root +from ....accuracy_core import GSM8K, LlmapiAccuracyTestHarness, assert_acceptance_length_for_llm +from ..._env import scoped_env_overrides + +# The target refuses other GPUs at construction: certification is per GPU architecture. +skip_not_sm100 = pytest.mark.skipif( + get_sm_version() != 100, reason="the glm5_next FP8 target is certified on sm_100 only" +) + + +class TestModelingV2Glm53FlashFp8Sm10xTp4(LlmapiAccuracyTestHarness): + """GLM-5.3-Flash block-FP8 / sm_100 / tp4, without speculation and with one-model MTP-3. + + The MTP variant also gates the acceptance length: a miscomputed draft layer only slows + decoding (every draft is rejected), so accuracy alone cannot see it. The acceptance + anchor is the built-in model's on the same checkpoint and MTP layer. + """ + + MODEL_NAME = "zai-org/GLM-5.3-Flash" + MODEL_PATH = f"{llm_models_root()}/GLM-5.3-Flash" + ACCEPTANCE_KEY = "ModelingV2Glm53FlashFp8Sm10xTp4::mtp3" + + @staticmethod + def _llm_kwargs(**extra) -> dict: + return dict( + tensor_parallel_size=4, + moe_tensor_parallel_size=4, + moe_expert_parallel_size=1, + kv_cache_config=KvCacheConfig( + dtype="fp8", free_gpu_memory_fraction=0.5, enable_block_reuse=False + ), + max_batch_size=64, + max_num_tokens=16384, + max_seq_len=8192, + cuda_graph_config=CudaGraphConfig(max_batch_size=64, enable_padding=True), + enable_chunked_prefill=True, + disable_mm_encoder=True, + modeling_v2="require", + **extra, + ) + + # Each engine also checks, on its first forward, that the metadata fields and cache members + # the target reads exist (TRTLLM_MODELING_V2_VALIDATE, passed to the ranks). + + @skip_not_sm100 + @pytest.mark.skip_less_device(4) + def test_gsm8k(self): + with ( + scoped_env_overrides({MODELING_V2_VALIDATE_ENV: "1"}) as rank_env, + LLM(self.MODEL_PATH, env_overrides=rank_env, **self._llm_kwargs()) as llm, + ): + GSM8K(self.MODEL_NAME).evaluate(llm) + self._check_one_token_prompt(llm) + + @staticmethod + def _check_one_token_prompt(llm: LLM) -> None: + """A one-token prompt runs as a one-row prefill step from a zero KDA state. + + Its output must not depend on which request held its slot before. Without speculation + the engine may run a final one-token context as a decode-graph row, which would skip the + slot reset the KDA prefill does; the target opts out of that, so the prompt takes the + one-row prefill path in both tests. + """ + params = SamplingParams(max_tokens=16, temperature=0) + token = llm.tokenizer.encode("Hello", add_special_tokens=False)[:1] + fresh = llm.generate(token, params).outputs[0].token_ids + llm.generate("Write a long story about a lighthouse keeper. " * 20, params) + after = llm.generate(token, params).outputs[0].token_ids + assert after == fresh, ( + f"one-token prompt changed with the slot's history: {fresh} -> {after}" + ) + + @skip_not_sm100 + @pytest.mark.skip_less_device(4) + def test_gsm8k_mtp3(self): + with ( + scoped_env_overrides({MODELING_V2_VALIDATE_ENV: "1"}) as rank_env, + LLM( + self.MODEL_PATH, + env_overrides=rank_env, + **self._llm_kwargs( + speculative_config=MTPDecodingConfig(max_draft_len=3), + max_stats_len=-1, + enable_iter_perf_stats=True, + ), + ) as llm, + ): + GSM8K(self.MODEL_NAME).evaluate(llm) + assert_acceptance_length_for_llm(self.ACCEPTANCE_KEY, llm) + self._check_one_token_prompt(llm) diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml index 756a84981b56..3bc910f3593b 100644 --- a/tests/integration/defs/accuracy/references/acceptance_length.yaml +++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml @@ -143,6 +143,13 @@ TestGLM53FlashFP8::test_mtp: TestModelingV2Glm53FlashNvfp4Sm100DefaultMtpFused::test_gsm8k_mtp: ref_al: 3.396 min_al: 3.226 +# GLM-5.3-Flash FP8, the modeling_v2 sm_10x tp4 target: the checkpoint's MTP layer chained for +# 3 drafts, measured on this target on 4x B200 (TP4, FP8 KV cache), GSM8K under the harness +# protocol (the built-in model measured 3.382 in the same setting). min_al is 95% of the +# measured value. +ModelingV2Glm53FlashFp8Sm10xTp4::mtp3: + ref_al: 3.380 + min_al: 3.211 # Measured on B200, 1P1D TP1 with target and draft KV in one manager. # min_al is 95% of the measured value. TestQwen3_8BDFlash::test_gsm8k_1p1d_unified_draft_kv: diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index cb45c57b2556..10088cde158c 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -197,6 +197,9 @@ l0_b200: - unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_persistent_sm_reserve.py - unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_rmsnorm.py - unittest/_torch/modeling_v2/gemm/test_modeling_v2_cublas_mm.py + # GLM-5.3-Flash FP8 modeling_v2 target (sm_100): its operations. The 4-rank boundary matrix runs in + # l0_dgx_b200. + - unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4 --ignore=unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/comm/test_modeling_v2_mhc_ar_boundary.py # KDA runtime state-pool and mixed-batch core coverage. - unittest/_torch/modules/kimi_kda/test_kimi_kda_bf16_state_pool.py - unittest/_torch/modules/kimi_kda/test_kda_mixed_batch_core_layout.py diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index f3806c6a4153..f81aac780a9c 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -120,6 +120,11 @@ l0_dgx_b200: # MTP-fused target with MTP-Eagle. - accuracy/modeling_v2/glm5_next/glm_5_3_flash_nvfp4__sm_100__default/test_glm5_next.py::TestModelingV2Glm53FlashNvfp4Sm100Default::test_gsm8k_mtp_vanilla - accuracy/modeling_v2/glm5_next/glm_5_3_flash_nvfp4__sm_100__default_mtp_fused/test_glm5_next.py::TestModelingV2Glm53FlashNvfp4Sm100DefaultMtpFused::test_gsm8k_mtp + # GLM-5.3-Flash FP8 modeling_v2 target: the fused all-reduce + mHC boundary, one 4-rank matrix + # (it skips with fewer than 4 GPUs, so the l0_b200 directory entry cannot run it), and the whole + # model with MTP-3: GSM8K, its acceptance length and the metadata contract. + - unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/comm/test_modeling_v2_mhc_ar_boundary.py + - accuracy/modeling_v2/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/test_accuracy.py::TestModelingV2Glm53FlashFp8Sm10xTp4::test_gsm8k_mtp3 # Exercise VisualGen worker-failure containment with live CUDA allocations # spread across four local processes. - unittest/_torch/visual_gen/test_executor_lifecycle_multi_gpu.py::test_sigkill_one_worker_contains_real_multi_gpu_group @@ -314,6 +319,8 @@ l0_dgx_b200: tests: # GLM-5.3-Flash modeling_v2 split target without speculation (SM100 only). - accuracy/modeling_v2/glm5_next/glm_5_3_flash_nvfp4__sm_100__default/test_glm5_next.py::TestModelingV2Glm53FlashNvfp4Sm100Default::test_gsm8k + # GLM-5.3-Flash FP8 modeling_v2 target: GSM8K without speculation. + - accuracy/modeling_v2/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/test_accuracy.py::TestModelingV2Glm53FlashFp8Sm10xTp4::test_gsm8k - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True-kv_dtype=fp8] - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True-kv_dtype=nvfp4] - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True-fuse_qkv_index_projection=False-eval_mode=inferencex-kv_dtype=fp8] diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/activation/test_modeling_v2_glm53_fp8_silu_and_mul.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/activation/test_modeling_v2_glm53_fp8_silu_and_mul.py new file mode 100644 index 000000000000..64bfde6d4d7d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/activation/test_modeling_v2_glm53_fp8_silu_and_mul.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the silu_and_mul operation of the glm5_next FP8 target. + +The entry owns what is certified. This file turns a cell's spec into tensors, shows that the +cell data tell the clamp and the gate/up order apart, and drives the inputs the guard exists +to refuse. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.activation.silu_and_mul import ( # noqa: E501 + GLM_SWIGLU_LIMIT, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.activation.silu_and_mul import ( # noqa: E501 + silu_and_mul as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +pytestmark = pytest.mark.skipif( + get_sm_version() != 100, reason="silu_and_mul is certified on sm_100" +) + +DEV = "cuda" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """torch's bf16 band: rtol 1.6e-2 (about four bf16 ulp), atol 1e-5. + + Kernel and reference evaluate the same fp32 expression on identical bf16 inputs and round + once. They differ only in the exponential inside the sigmoid (Triton's approximate exp + against torch's), a few fp32 ulp, which can move a bf16 rounding by one ulp (0.4% + relative) and no further. Where a very negative gate drives the product toward zero the + two sit far inside the 1e-5 absolute floor. Each wrong reading this entry guards against + -- no clamp, a symmetric gate clamp, the halves swapped -- misses by orders of magnitude on + the elements it touches (`test_cell_data_tell_the_clamp_and_the_halves_apart`). + """ + assert out.dtype == ref.dtype and out.shape == ref.shape, (out.dtype, ref.dtype) + torch.testing.assert_close(out, ref) + + +def _build(spec: dict, seed: int) -> torch.Tensor: + """FC1 outputs at the scale a GLM layer's look like: std 8, so the limit of 10 clamps + about 11% of the gates (upper side) and 21% of the ups (both sides).""" + g = torch.Generator(device=DEV).manual_seed(seed) + x = torch.randn(spec["rows"], spec["cols"], generator=g, device=DEV) * 8.0 + return x.to(torch.bfloat16) + + +# The shapes the target calls this operation with, on top of the operation's own `CELLS`. +# +# GLM-5.3-Flash at TP4: the three dense layers' MLP (intermediate 12288 / 4 = 3072, so 2 x 3072 +# wide) at every step, and the MoE layers' shared expert (2048 / 4 = 512) on steps the +# routed-expert decode path does not take (more than 4 rows). Generation steps are MTP-3 +# verifies of 4 rows per request, up to 64 requests; a 13,312-row prefill runs both on every +# gathered row. +CELLS: tuple[Cell, ...] = op.CELLS + ( + Cell( + why="dense MLP, MTP-3 verify of one request: 4 rows of 2 x 3072", + spec=dict(rows=4, cols=6144, swiglu_limit=GLM_SWIGLU_LIMIT), + ), + Cell( + why="dense MLP, MTP-3 verify of 64 requests: 256 rows of 2 x 3072", + spec=dict(rows=256, cols=6144, swiglu_limit=GLM_SWIGLU_LIMIT), + ), + Cell( + why="dense MLP over a 13,312-row prefill step", + spec=dict(rows=13312, cols=6144, swiglu_limit=GLM_SWIGLU_LIMIT), + ), + Cell( + why="shared expert, the smallest step past the 4-row decode path: 5 rows of 2 x 512", + spec=dict(rows=5, cols=1024, swiglu_limit=GLM_SWIGLU_LIMIT), + ), + Cell( + why="shared expert, MTP-3 verify of 64 requests: 256 rows of 2 x 512", + spec=dict(rows=256, cols=1024, swiglu_limit=GLM_SWIGLU_LIMIT), + ), + Cell( + why="shared expert over a 13,312-row prefill step", + spec=dict(rows=13312, cols=1024, swiglu_limit=GLM_SWIGLU_LIMIT), + ), +) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture + certifying(op) + x = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + limit = cell.spec["swiglu_limit"] + out = op(x, swiglu_limit=limit) + assert out.shape == (x.shape[0], x.shape[1] // 2) and out.dtype == x.dtype + assert out.is_contiguous() and out.data_ptr() != x.data_ptr() + _compare(out, op.reference(x, swiglu_limit=limit)) + + +def test_cell_data_tell_the_clamp_and_the_halves_apart() -> None: + """The cells certify the clamp and the half order only if their data can see them. + + At the cell scale each wrong reading misses: no clamp, the gate clamped on both sides + (it is clamped on the upper side only), and the halves swapped ([up | gate] instead of + [gate | up]). + """ + spec = dict(rows=256, cols=1024, swiglu_limit=GLM_SWIGLU_LIMIT) + x = _build(spec, seed=7) + out = op(x, swiglu_limit=GLM_SWIGLU_LIMIT) + _compare(out, op.reference(x, swiglu_limit=GLM_SWIGLU_LIMIT)) + + d = x.shape[1] // 2 + gate, up = x[:, :d].float(), x[:, d:].float() + lim = GLM_SWIGLU_LIMIT + up_c = up.clamp(-lim, lim) + symmetric = gate.clamp(-lim, lim) + wrong = { + "no clamp": op.reference(x), + "symmetric gate clamp": (symmetric * torch.sigmoid(symmetric) * up_c).to(x.dtype), + "halves swapped": op.reference(torch.cat([x[:, d:], x[:, :d]], dim=1), swiglu_limit=lim), + } + for name, ref in wrong.items(): + with pytest.raises(AssertionError): + _compare(out, ref) + print(f" {name}: rejected") + + +def test_a_row_strided_view_is_read_correctly() -> None: + """Rows are read through `x.stride(0)`, so a column slice of a wider buffer that keeps the + last dim dense needs no copy -- the guard admits it.""" + buf = _build(dict(rows=8, cols=8192), seed=11) + x = buf[:, :6144] + assert x.stride(0) == 8192 and x.stride(1) == 1 + with validating(op): + out = op(x, swiglu_limit=GLM_SWIGLU_LIMIT) + _compare(out, op.reference(x.contiguous(), swiglu_limit=GLM_SWIGLU_LIMIT)) + + +def test_guard_refuses_a_strided_last_dim() -> None: + """The hazard: a strided last dim is read as dense and the answer is silently wrong.""" + buf = _build(dict(rows=8, cols=2048), seed=13) + x = buf[:, ::2] + assert x.stride(-1) == 2 + unguarded = op(x, swiglu_limit=GLM_SWIGLU_LIMIT) + with pytest.raises(AssertionError): + _compare(unguarded, op.reference(x, swiglu_limit=GLM_SWIGLU_LIMIT)) + with pytest.raises(AssertionError, match="last dim must be contiguous"): + with validating(op): + op(x, swiglu_limit=GLM_SWIGLU_LIMIT) + + +def test_guard_refuses_a_scale_without_a_dtype() -> None: + x = _build(dict(rows=4, cols=1024), seed=17) + scale = torch.ones(1, dtype=torch.float32, device=DEV) + with pytest.raises(AssertionError, match="scale without dtype"): + with validating(op): + op(x, scale=scale, swiglu_limit=GLM_SWIGLU_LIMIT) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_compact_sparse_rows.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_compact_sparse_rows.py new file mode 100644 index 000000000000..04a0d962584c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_compact_sparse_rows.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for the glm5_next fp8 tp4 entry compact_sparse_rows.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.attention.compact_sparse_rows import ( # noqa: E501 + compact_sparse_rows, +) +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() != 100, reason="the glm5_next fp8 tp4 entries are certified on sm_100" +) + +DEVICE = "cuda" +SELECT_K, KPOOL = 512, 4 + + +def _expanded(spec: dict, seed: int) -> torch.Tensor: + """Rows shaped like ``kpool_expand``'s output, from each row's visible length (the cell's + ``kv`` for the last row, one less per row before it): the members of its selected pools, + its tail, ``-1`` padding. 5% of the entries are knocked out to ``-1`` so that every row has + holes to close, and the last row selects nothing.""" + g = torch.Generator(device=DEVICE).manual_seed(seed) + rows, width = spec["rows"], spec["width"] + visible = (spec["kv"] - torch.arange(rows, device=DEVICE).flip(0)).clamp(min=0) + pools = (visible // KPOOL).clamp(max=SELECT_K) + col = torch.arange(width, device=DEVICE) + members = (col < SELECT_K * KPOOL) & (col // KPOOL < pools[:, None]) + tail = col - SELECT_K * KPOOL + tail = (tail >= 0) & (tail < (visible % KPOOL)[:, None]) + ids = torch.randint(0, 1 << 24, (rows, width), generator=g, device=DEVICE, dtype=torch.int32) + holes = torch.rand(rows, width, generator=g, device=DEVICE) < 0.05 + indices = torch.where((members | tail) & ~holes, ids, -1) + indices[-1] = -1 + return indices + + +@pytest.mark.parametrize( + "cell", compact_sparse_rows.CELLS, ids=[c.why[:44] for c in compact_sparse_rows.CELLS] +) +def test_compact_matches_reference(cell, certifying): # noqa: F811 + """Packed ids, lengths and the nonempty mask equal the reference's exactly: the op moves + integers, so there is no tolerance. The last row, which selects nothing, gets the dummy + id 0, length 1 and nonempty False.""" + certifying(compact_sparse_rows) + indices = _expanded(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + out = compact_sparse_rows(indices) + ref = compact_sparse_rows.reference(indices) + for got, want in zip(out, ref): + assert got.dtype == want.dtype and got.shape == want.shape + assert torch.equal(got, want) + packed, lengths, nonempty = out + assert int(packed[-1, 0]) == 0 and int(lengths[-1]) == 1 and not bool(nonempty[-1]) + + +def test_guard_refuses_column_strided_ids(): + """A column-strided view is read as if its columns were adjacent: the op would pack + entries of the underlying buffer the view does not hold.""" + buf = torch.randint(-1, 100, (4, 2 * 2112), device=DEVICE, dtype=torch.int32) + with pytest.raises(AssertionError, match="unit-stride"): + with validating(compact_sparse_rows): + compact_sparse_rows(buf[:, ::2]) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_head_major.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_head_major.py new file mode 100644 index 000000000000..c57b9c4457e9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_head_major.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for the glm5_next fp8 tp4 head-major FP8 cast entries.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.attention.head_major import ( # noqa: E501 + quantize_1x128, + saturate_fp8, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 100, reason="the glm5_next fp8 tp4 entries are certified on sm_100" +) + + +def _input(spec: dict, seed: int, scale: float) -> torch.Tensor: + g = torch.Generator(device="cuda").manual_seed(seed) + x = torch.randn(spec["h"], spec["r"], spec["d"], generator=g, device="cuda") * scale + return x.to(torch.bfloat16) + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.uint8) + + +@pytest.mark.parametrize("cell", saturate_fp8.CELLS, ids=[c.why[:44] for c in saturate_fp8.CELLS]) +def test_saturate_matches_bitwise(cell, certifying): # noqa: F811 + """Values past +-448 must saturate, not become NaN; the rest round as torch's cast does.""" + certifying(saturate_fp8) + x = _input(cell.spec, zlib.crc32(cell.why.encode()) % 2**31, scale=300.0) + saturate_fp8.is_valid(x) + assert torch.equal(_bits(saturate_fp8(x)), _bits(saturate_fp8.reference(x))) + + +@pytest.mark.parametrize( + "cell", quantize_1x128.CELLS, ids=[c.why[:44] for c in quantize_1x128.CELLS] +) +def test_quantize_matches_bitwise(cell, certifying): # noqa: F811 + """Same codes and scales as ``fp8_quantize_1x128`` on the transposed copy.""" + certifying(quantize_1x128) + x = _input(cell.spec, zlib.crc32(cell.why.encode()) % 2**31, scale=1.0) + x[0, :3] = 0 # all-zero blocks take the 1e-10 amax floor + quantize_1x128.is_valid(x) + (q, s), (q_ref, s_ref) = quantize_1x128(x), quantize_1x128.reference(x) + assert torch.equal(_bits(q), _bits(q_ref)) + assert torch.equal(s, s_ref) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_indexer_topk.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_indexer_topk.py new file mode 100644 index 000000000000..72bfca13edd9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_indexer_topk.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for the glm5_next fp8 tp4 entry indexer_topk.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.attention.indexer_topk import ( # noqa: E501 + indexer_topk, +) +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() != 100, reason="the glm5_next fp8 tp4 entries are certified on sm_100" +) + +DEVICE = "cuda" +KPOOL, SELECT_K = 4, 512 + + +def _args(spec: dict, seed: int): + """Scores and candidate counts the way the target hands them over. + + Context cells: rows ``lo .. hi - 1`` of a prompt see ``row + 1`` positions; the scores are + a view of DeepGEMM's row-padded logits. Generation cells: request ``b``'s KV length after + the step is ``kv - b % 4`` and its rows see ``kv_b - T + 1 .. kv_b`` positions, over rows + a full block table wide. Every column a row must not read (at + or past its count, and the row padding) holds +inf, so a read past the count wins the + selection. ``levels`` rounds the scores to quarter steps: the k-th score ties many times. + """ + g = torch.Generator(device=DEVICE).manual_seed(seed) + if "rows" in spec: + lo, hi = spec["rows"] + visible = torch.arange(lo + 1, hi + 1, device=DEVICE) + else: + kv = spec["kv"] - torch.arange(spec["requests"], device=DEVICE) % 4 + steps = torch.arange(spec["tokens"], device=DEVICE) + visible = (kv[:, None] - spec["tokens"] + 1 + steps).reshape(-1) + num_cand = (visible // KPOOL).to(torch.int32) + rows, width = visible.shape[0], spec["width"] + buf = torch.randn(rows, width + spec.get("row_pad", 0), generator=g, device=DEVICE) + if "levels" in spec: + buf = torch.round(buf * spec["levels"]) / spec["levels"] + col = torch.arange(buf.shape[1], device=DEVICE) + buf = torch.where(col < num_cand[:, None].long(), buf, float("inf")) + scores = buf[:, :width] + selected = torch.full((rows, SELECT_K), 12345, dtype=torch.int32, device=DEVICE) + return scores, num_cand, selected + + +def _check(out, ref, scores, num_cand): + """Rows with at most ``SELECT_K`` candidates: exactly the reference (``0 .. n - 1``, then + -1). Longer rows: ``SELECT_K`` distinct candidates whose scores, sorted, equal the + reference's bit for bit (true of every valid top-k, whatever the tie rule), and the same + pools as the reference wherever the k-th score is not tied.""" + short = num_cand.long() <= SELECT_K + assert torch.equal(out[short], ref[short]) + long_rows = (~short).nonzero().flatten() + if long_rows.numel() == 0: + return + got = out[long_rows].long().sort(dim=1).values + want = ref[long_rows].long() + count = num_cand[long_rows, None].long() + assert bool(((got >= 0) & (got < count)).all()), "a pool outside the row's candidates" + assert bool((got.diff(dim=1) > 0).all()), "a pool selected twice" + got_scores = torch.gather(scores[long_rows], 1, got).sort(dim=1).values + want_scores = torch.gather(scores[long_rows], 1, want).sort(dim=1).values + assert torch.equal(got_scores, want_scores) + # Without a tie at the k-th score the selected set is unique. + col = torch.arange(scores.shape[1], device=DEVICE) + candidates = torch.where(col < count, scores[long_rows], float("-inf")) + next_best = candidates.topk(SELECT_K + 1, dim=1).values[:, -1] + untied = want_scores[:, 0] > next_best + assert torch.equal(got[untied], want[untied]) + + +@pytest.mark.parametrize("cell", indexer_topk.CELLS, ids=[c.why[:44] for c in indexer_topk.CELLS]) +def test_indexer_topk_matches_reference(cell, certifying): # noqa: F811 + """See ``_check``: the kernel's choice among tied pools and its order are unspecified, so + the comparison is of the selected scores and, where nothing ties, of the pool sets.""" + certifying(indexer_topk) + scores, num_cand, selected = _args(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + out = indexer_topk(scores, num_cand, selected, select_k=SELECT_K) + assert out is selected + ref = indexer_topk.reference(scores, num_cand, selected, select_k=SELECT_K) + _check(out, ref, scores, num_cand) + + +def test_guard_refuses_counts_past_the_row(): + """A candidate count past the row reads the next row's scores as this row's.""" + scores, num_cand, selected = _args(dict(requests=1, tokens=4, kv=4096, width=1000), seed=0) + with pytest.raises(AssertionError, match="candidate counts"): + with validating(indexer_topk): + indexer_topk(scores, num_cand, selected, select_k=SELECT_K) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_kpool_expand.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_kpool_expand.py new file mode 100644 index 000000000000..d1c283ad09cb --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_kpool_expand.py @@ -0,0 +1,122 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for the glm5_next fp8 tp4 entry kpool_expand.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.attention.kpool_expand import ( # noqa: E501 + kpool_expand, +) +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() != 100, reason="the glm5_next fp8 tp4 entries are certified on sm_100" +) + +DEVICE = "cuda" +TPB, KPOOL, SELECT_K, OUT_WIDTH, PAGES = 32, 4, 512, 2112, 4100 +# A coalesced latent pool: this layer's 32 rows of a slot start at row 32 of 96. +BASE_ROW, ROWS_PER_SLOT = 32, 96 + + +def _schedule(spec: dict): + """``(visible int64 [N], request_ids int32 [N] or None, tables)`` of a cell. + + Generation cells: request ``b``'s KV length after the step is ``kv - b % 4``. With ``T > 1`` + its rows are positions ``kv_b - T .. kv_b - 1`` and read table ``1 + b`` by request id (row + 0 is a context request's table, as in a mixed step); with ``T = 1`` row ``b`` reads table + ``b`` of the generation tables, without request ids. Context cells: the packed rows of each + ``(cached, new)`` request, restricted to ``rows``, reading their request's table. + """ + if "context" in spec: + visible, request_ids = [], [] + for i, (cached, new) in enumerate(spec["context"]): + visible.append(torch.arange(cached + 1, cached + new + 1)) + request_ids.append(torch.full((new,), i, dtype=torch.int32)) + lo, hi = spec["rows"] + visible = torch.cat(visible)[lo:hi].to(DEVICE) + return visible, torch.cat(request_ids)[lo:hi].to(DEVICE), len(spec["context"]) + requests, tokens = spec["requests"], spec["tokens"] + kv = spec["kv"] - torch.arange(requests, device=DEVICE) % 4 + if tokens == 1: + return kv, None, requests + visible = (kv[:, None] - tokens + 1 + torch.arange(tokens, device=DEVICE)).reshape(-1) + request_ids = (1 + torch.arange(requests, device=DEVICE, dtype=torch.int32)).repeat_interleave( + tokens + ) + return visible, request_ids, 1 + requests + + +def _top_k_like(visible: torch.Tensor, g: torch.Generator) -> torch.Tensor: + """Selections shaped like the decode top-k's: a row with at most ``SELECT_K`` candidate + pools lists them in order then ``-1``; a longer row lists ``SELECT_K`` distinct candidates + in random order.""" + num_cand = visible // KPOOL + width = max(int(num_cand.max()), SELECT_K) + keys = torch.rand(visible.shape[0], width, generator=g, device=DEVICE) + keys = torch.where(torch.arange(width, device=DEVICE) < num_cand[:, None], keys, -1.0) + selected = keys.topk(SELECT_K, dim=1).indices.to(torch.int32) + ident = torch.arange(SELECT_K, device=DEVICE, dtype=torch.int32) + short = num_cand <= SELECT_K + selected[short] = torch.where(ident < num_cand[short, None], ident, -1).to(torch.int32) + return selected + + +def _args(spec: dict, seed: int): + g = torch.Generator(device=DEVICE).manual_seed(seed) + visible, request_ids, num_tables = _schedule(spec) + pages = spec.get("pages", PAGES) + tables = torch.randint(0, 1 << 20, (num_tables, pages), generator=g, device=DEVICE) + selected = _top_k_like(visible, g) + # The kernel's own validity mask: a -1 and an id past the candidates in a long row. + long_rows = (visible // KPOOL > SELECT_K).nonzero().flatten() + if long_rows.numel(): + r = int(long_rows[0]) + selected[r, 5] = int(visible[r] // KPOOL) + 7 + selected[r, 9] = -1 + args = (selected, visible, tables, TPB) + kwargs = dict( + base_row=BASE_ROW, + rows_per_slot=ROWS_PER_SLOT, + kpool=KPOOL, + out_width=OUT_WIDTH, + request_ids=request_ids, + ) + return args, kwargs + + +@pytest.mark.parametrize("cell", kpool_expand.CELLS, ids=[c.why[:44] for c in kpool_expand.CELLS]) +def test_kpool_expand_matches_reference(cell, certifying): # noqa: F811 + """The row ids equal the reference's exactly (integer arithmetic): members of every valid + selected pool, the tail, -1 for invalid pools and the padding.""" + certifying(kpool_expand) + args, kwargs = _args(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + out = kpool_expand(*args, **kwargs) + ref = kpool_expand.reference(*args, **kwargs) + assert out.shape == ref.shape == (args[0].shape[0], OUT_WIDTH) + assert torch.equal(out, ref) + + +def test_guard_refuses_visible_lengths_past_the_table(): + """A visible position past a table's last page reads the next table's pages as its own.""" + selected = torch.full((1, SELECT_K), -1, dtype=torch.int32, device=DEVICE) + tables = torch.zeros(2, 4, dtype=torch.int64, device=DEVICE) + kv_lens = torch.tensor([4 * TPB + 3], device=DEVICE) + with pytest.raises(AssertionError, match="exceeds"): + with validating(kpool_expand): + kpool_expand( + selected, + kv_lens, + tables, + TPB, + base_row=0, + rows_per_slot=TPB, + kpool=KPOOL, + out_width=OUT_WIDTH, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_kpool_score.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_kpool_score.py new file mode 100644 index 000000000000..fd30273b1b88 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_kpool_score.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for the glm5_next fp8 tp4 entry kpool_score.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.attention.kpool_score import ( # noqa: E501 + FP32_MIN, + kpool_score, +) +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() != 100, reason="the glm5_next fp8 tp4 entries are certified on sm_100" +) + +DEVICE = "cuda" +TPB, KPOOL, HD, HEADS, PAGES, SLOTS = 32, 4, 128, 32, 4100, 1024 +Q_SCALE, W_SCALE = HD**-0.5, HEADS**-0.5 + + +def _schedule(spec: dict) -> tuple[torch.Tensor, int]: + """``(visible int64 [N], tables)`` of a cell. + + Request ``b``'s KV length after the step is ``kv - b % 4``; with ``T > 1`` its rows are + positions ``kv_b - T .. kv_b - 1`` (visible ``position + 1``) of generation table ``b``, + which follows one context request's table. + """ + requests, tokens = spec["requests"], spec["tokens"] + kv = spec["kv"] - torch.arange(requests, device=DEVICE) % 4 + if tokens == 1: + return kv, requests + visible = (kv[:, None] - tokens + 1 + torch.arange(tokens, device=DEVICE)).reshape(-1) + return visible, 1 + requests + + +def _args(spec: dict, seed: int) -> tuple[tuple, dict]: + """The op's arguments the way the target passes them: per-row copies of the generation + tables (``GenerationRows.tables``, int64). The pool is a coalesced view: each slot's rows sit + in the second of two subpages, with padded rows.""" + g = torch.Generator(device=DEVICE).manual_seed(seed) + visible, num_tables = _schedule(spec) + pages, slots = spec.get("pages", PAGES), spec.get("slots", SLOTS) + used = -(-int(visible.max()) // TPB) + tables = torch.zeros(num_tables, pages, dtype=torch.int64, device=DEVICE) + for t in range(num_tables): + tables[t, :used] = torch.randperm(slots, generator=g, device=DEVICE)[:used] + storage = torch.randn(slots, 2, TPB, 3 * HD + 64, generator=g, device=DEVICE) + pool = storage.to(torch.bfloat16)[:, 1, :, : 3 * HD] + n = visible.shape[0] + q = torch.randn(n, HEADS, HD, generator=g, device=DEVICE).to(torch.bfloat16) + # The head weights are the tail of the indexer's [k | gate | weights] GEMM output. + projection = torch.randn(n, 2 * HD + HEADS, generator=g, device=DEVICE).to(torch.bfloat16) + weights = projection[:, 2 * HD :] + kwargs = dict( + num_pools_max=-(-pages * TPB // KPOOL), + head_dim=HD, + kpool=KPOOL, + q_scale=Q_SCALE, + w_scale=W_SCALE, + ) + if spec["tokens"] > 1: + tables = tables[1:].repeat_interleave(spec["tokens"], dim=0) + return (q, weights, pool, tables, visible, TPB), kwargs + + +def _check(out: torch.Tensor, args: tuple, kwargs: dict) -> None: + """Every visible score within ``2**-13 * bound`` of the reference, where ``bound = q_scale + * sum_h |w_h w_scale| sum_d |q_hd k_d|`` (the reference on absolute values). Both sides + sum exact BF16 products in FP32, 128 per head and 32 heads, so each is within ``~161 * + 2**-24 * bound`` of the exact score; ``2**-13`` leaves eight times that for the + tensor-core accumulation order. Every invisible column is the FP32 minimum.""" + ref = kpool_score.reference(*args, **kwargs) + q, weights, pool, tables, kv_lens, tpb = args + bound = kpool_score.reference( + q.abs(), weights.abs(), pool.abs(), tables, kv_lens, tpb, **kwargs + ) + assert out.shape == ref.shape and out.dtype == torch.float32 + col = torch.arange(out.shape[1], device=DEVICE) + visible = col[None, :] < (kv_lens.long() // KPOOL)[:, None] + err = (out - ref).abs()[visible] + worst = float((err / bound[visible].clamp_min(1e-30)).max()) if err.numel() else 0.0 + assert bool((err <= 2**-13 * bound[visible]).all()), f"worst error {worst:.3e} of the bound" + assert bool((out[~visible] == FP32_MIN).all()) + + +@pytest.mark.parametrize("cell", kpool_score.CELLS, ids=[c.why[:44] for c in kpool_score.CELLS]) +def test_kpool_score_matches_reference(cell, certifying): # noqa: F811 + """Visible pool scores within the accumulation band of the FP32 reference, every invisible + pool at the FP32 minimum; see ``_check`` for the band.""" + certifying(kpool_score) + args, kwargs = _args(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + out = kpool_score(*args, **kwargs) + _check(out, args, kwargs) + + +def test_guard_refuses_fp32_queries(): + """The dot reads the queries as BF16: an FP32 query would be rounded without a sign.""" + args, kwargs = _args(dict(requests=1, tokens=4, kv=600, pages=32, slots=64), seed=0) + q = args[0].float() + with pytest.raises(AssertionError, match="BF16"): + with validating(kpool_score): + kpool_score(q, *args[1:], **kwargs) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_kpool_score_fp8.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_kpool_score_fp8.py new file mode 100644 index 000000000000..cad5d4f3a73b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_kpool_score_fp8.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for the glm5_next fp8 tp4 entry kpool_score_fp8.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.attention.kpool_score_fp8 import ( # noqa: E501 + kpool_score_fp8, +) +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() != 100, reason="the glm5_next fp8 tp4 entries are certified on sm_100" +) + +DEVICE = "cuda" +TPB, KPOOL, HD, HEADS, PAGES, SLOTS = 32, 4, 128, 32, 4100, 1024 +Q_SCALE, W_SCALE = HD**-0.5, HEADS**-0.5 + + +def _on_e4m3_grid(shape, g: torch.Generator) -> torch.Tensor: + """BF16 rows of 128 that E4M3 holds exactly at the quantizer's per-row scale: each row is + E4M3 values times ``2**e`` (``e`` drawn per row) with an absolute maximum of exactly ``448 * + 2**e``, so ``amax / 448`` is a power of two and every value is a code times it.""" + x = torch.randn(*shape, generator=g, device=DEVICE) + x = (x / x.abs().amax(-1, keepdim=True) * 448).to(torch.float8_e4m3fn).float() + top = x.abs().argmax(-1, keepdim=True) + x.scatter_(-1, top, 448 * torch.gather(x, -1, top).sign()) + e = torch.randint(-12, -1, (*shape[:-1], 1), generator=g, device=DEVICE) + return (x * torch.exp2(e.float())).to(torch.bfloat16) + + +def _args(spec: dict, seed: int, exact: bool = True) -> tuple[tuple, dict]: + """Packed context rows of ``(cached, new)`` requests, restricted to ``rows``, with the + host bound ``max_kv_len`` the target passes: the largest visible length among those rows. + With ``exact`` the queries and every pool key are on the E4M3 grid (``_on_e4m3_grid``).""" + g = torch.Generator(device=DEVICE).manual_seed(seed) + visible, request_ids = [], [] + for i, (cached, new) in enumerate(spec["context"]): + visible.append(torch.arange(cached + 1, cached + new + 1)) + request_ids.append(torch.full((new,), i, dtype=torch.int32)) + lo, hi = spec["rows"] + visible = torch.cat(visible)[lo:hi].to(DEVICE) + request_ids = torch.cat(request_ids)[lo:hi].to(DEVICE) + pages, slots = spec.get("pages", PAGES), spec.get("slots", SLOTS) + max_kv_len = int(visible.max()) + tables = torch.zeros(len(spec["context"]), pages, dtype=torch.int64, device=DEVICE) + for t, (cached, new) in enumerate(spec["context"]): + used = -(-(cached + new) // TPB) + tables[t, :used] = torch.randperm(slots, generator=g, device=DEVICE)[:used] + storage = torch.randn(slots, 2, TPB, 3 * HD + 64, generator=g, device=DEVICE) + pool = storage.to(torch.bfloat16)[:, 1, :, : 3 * HD] + n = visible.shape[0] + if exact: + pool[..., 2 * HD :] = _on_e4m3_grid((slots, TPB, HD), g) + q = _on_e4m3_grid((n, HEADS, HD), g) + else: + q = torch.randn(n, HEADS, HD, generator=g, device=DEVICE).to(torch.bfloat16) + projection = torch.randn(n, 2 * HD + HEADS, generator=g, device=DEVICE).to(torch.bfloat16) + weights = projection[:, 2 * HD :] + kwargs = dict( + head_dim=HD, + kpool=KPOOL, + q_scale=Q_SCALE, + w_scale=W_SCALE, + request_ids=request_ids, + max_kv_len=max_kv_len, + ) + return (q, weights, pool, tables, visible, TPB), kwargs + + +def _visible(args, pools): + col = torch.arange(pools, device=DEVICE) + return col[None, :] < (args[4].long() // KPOOL)[:, None] + + +def _check_layout(out, args, kwargs): + """One table: DeepGEMM's row-padded logits as a strided view; several: the aligned copy.""" + n, tables = args[0].shape[0], args[3].shape[0] + pools = kwargs["max_kv_len"] // KPOOL + assert out.dtype == torch.float32 and out.stride(1) == 1 and out.shape[0] == n + if tables == 1: + assert out.shape[1] == pools and out.stride(0) >= pools and out.stride(0) % 4 == 0 + else: + assert out.shape[1] == max(1024, -(-pools // 64) * 64) and out.is_contiguous() + + +@pytest.mark.parametrize( + "cell", kpool_score_fp8.CELLS, ids=[c.why[:44] for c in kpool_score_fp8.CELLS] +) +def test_kpool_score_fp8_matches_reference_on_the_e4m3_grid(cell, certifying): # noqa: F811 + """Operands on the E4M3 grid quantize without loss (the query and key scales come out as + powers of two, up to one rounding of ``1/448``), so the FP8 scores equal the exact ones up + to FP32 accumulation: within ``2**-13 * bound``, ``bound`` the reference on absolute + values, as for ``kpool_score`` (exact products of 4-bit significands, 128 per head, 32 + heads). This pins the gather, the table ranges and every scale, not the rounding.""" + certifying(kpool_score_fp8) + args, kwargs = _args(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + out = kpool_score_fp8(*args, **kwargs) + ref = kpool_score_fp8.reference(*args, **kwargs) + _check_layout(out, args, kwargs) + q, weights, pool, tables, kv_lens, tpb = args + bound = kpool_score_fp8.reference( + q.abs(), weights.abs(), pool.abs(), tables, kv_lens, tpb, **kwargs + ) + pools = ref.shape[1] + visible = _visible(args, pools) + err = (out[:, :pools] - ref).abs()[visible] + worst = float((err / bound[visible].clamp_min(1e-30)).max()) + assert bool((err <= 2**-13 * bound[visible]).all()), f"worst error {worst:.3e} of the bound" + + +def test_kpool_score_fp8_rounding_stays_in_band(): + """General BF16 operands: the E4M3 rounding of the queries and keys moves each score. + + With per-row scale ``s = amax / 448`` an operand ``x`` rounds by at most ``h(x) = + max(2**-4 |x|, 2**-10 s)`` (half an E4M3 step, the subnormal step below ``2**-6 s``), + modelled as uniform: variance ``h**2 / 3``. A score's error is then a sum over heads and + dims of independent terms of variance ``(w'_h)**2 (h(q)**2 k**2 + q**2 h(k)**2) / 3`` + (``w'_h = w_h * w_scale * q_scale``; the relu only shrinks it). Every visible score must + sit within 8 of those standard deviations plus the accumulation band of the exact-grid + test; across the ~1.4M visible scores of this cell a Gaussian tail past 8 sigma is below 1e-8. + """ + cell = kpool_score_fp8.CELLS[0] + args, kwargs = _args(cell.spec, zlib.crc32(cell.why.encode()) % 2**31, exact=False) + out = kpool_score_fp8(*args, **kwargs) + ref = kpool_score_fp8.reference(*args, **kwargs) + q, weights, pool, tables, kv_lens, tpb = args + pools = ref.shape[1] + first = torch.arange(pools, device=DEVICE) * KPOOL + keys = pool[:, :, 2 * HD :][tables[:, first // TPB], first % TPB].float() # [T, P, HD] + + def half_step(x): + scale = x.abs().amax(-1, keepdim=True) / 448 + return torch.maximum(2**-4 * x.abs(), 2**-10 * scale) + + qf = q.float() + hq, hk = half_step(qf), half_step(keys) + head_w = (weights.float() * W_SCALE * Q_SCALE) ** 2 # [N, H] + table = kwargs["request_ids"].long() + var = torch.empty_like(ref) + for r0 in range(0, q.shape[0], 64): + rows = slice(r0, min(r0 + 64, q.shape[0])) + k2, hk2 = keys[table[rows]] ** 2, hk[table[rows]] ** 2 # [c, P, HD] + per_head = torch.einsum("rhd,rpd->rhp", hq[rows] ** 2, k2) + torch.einsum( + "rhd,rpd->rhp", qf[rows] ** 2, hk2 + ) + var[rows] = (per_head * head_w[rows, :, None]).sum(dim=1) / 3 + bound = kpool_score_fp8.reference( + q.abs(), weights.abs(), pool.abs(), tables, kv_lens, tpb, **kwargs + ) + visible = _visible(args, pools) + err = (out[:, :pools] - ref).abs() + band = 8 * var.sqrt() + 2**-13 * bound + worst = float((err / band)[visible].max()) + assert bool((err <= band)[visible].all()), f"worst error {worst:.3f} of the band" + + +def test_guard_refuses_lengths_past_max_kv_len(): + """Row r's key range is ``request_ids[r] * pools`` plus its pool count: a visible length + past ``max_kv_len`` runs into the next table's keys.""" + spec = dict(context=((0, 300), (0, 300)), rows=(0, 600), pages=16, slots=64) + args, kwargs = _args(spec, seed=0) + kwargs["max_kv_len"] = 200 + with pytest.raises(AssertionError, match="max_kv_len"): + with validating(kpool_score_fp8): + kpool_score_fp8(*args, **kwargs) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_kpool_update.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_kpool_update.py new file mode 100644 index 000000000000..e55d34fda807 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_kpool_update.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for the glm5_next fp8 tp4 entry kpool_update.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.attention.kpool_update import ( # noqa: E501 + kpool_update, +) +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() != 100, reason="the glm5_next fp8 tp4 entries are certified on sm_100" +) + +DEVICE = "cuda" +TPB, KPOOL, HD, PAGES = 32, 4, 128, 4100 + + +def _rows(spec: dict): + """``(positions, request_ids, pages each table needs)`` of a cell. + + Generation cells: request ``b``'s KV length after the step is ``kv - b % 4``; its ``T`` + rows are positions ``kv_b - T .. kv_b - 1`` (``[B, T]``, or ``[B]`` when ``T = 1``) of + generation table ``b``, whose pools reach back to the first member of the first row's + pool. Context cells: the pool-final positions each ``(cached, new)`` request completes in + the step (one row per pool, with request ids), as the target's context schedule + (``ContextRows.final_positions``) lists them; their pools span every page of the request. + """ + if "context" in spec: + positions, request_ids, pages = [], [], [] + for i, (cached, new) in enumerate(spec["context"]): + first_final = cached + (-(cached + 1)) % KPOOL + finals = list(range(first_final, cached + new, KPOOL)) + positions += finals + request_ids += [i] * len(finals) + pages.append(range(0, -(-(cached + new) // TPB))) + return ( + torch.tensor(positions, dtype=torch.int64, device=DEVICE), + torch.tensor(request_ids, dtype=torch.int32, device=DEVICE), + pages, + ) + requests, tokens = spec["requests"], spec["tokens"] + kv = spec["kv"] - torch.arange(requests, device=DEVICE) % 4 + positions = (kv[:, None] - tokens + torch.arange(tokens, device=DEVICE)).squeeze(1) + pages = [ + range((length - tokens) // KPOOL * KPOOL // TPB, (length - 1) // TPB + 1) + for length in kv.tolist() + ] + return positions, None, pages + + +def _args(spec: dict, seed: int): + """The op's arguments and the storage behind the pool view. + + Every page a table's rows reach gets its own slot (the other pages point at slot 0, which + no row reaches). The pool is a coalesced view: the second of two subpages of every slot, + rows padded by 64 columns, as a layer's indexer pool sits in the cache's pages. + Generation tables follow one context request's table, so the op reads a row slice of the + step's tables, like the target's generation slice. + """ + g = torch.Generator(device=DEVICE).manual_seed(seed) + positions, request_ids, pages = _rows(spec) + lead = 0 if request_ids is not None else 1 + tables = torch.zeros(lead + len(pages), spec.get("pages", PAGES), dtype=torch.int64) + slots = 1 + sum(len(p) for p in pages) + order = (1 + torch.randperm(slots - 1, generator=torch.Generator().manual_seed(seed))).tolist() + for t, table_pages in enumerate(pages, start=lead): + for page in table_pages: + tables[t, page] = order.pop() + tables = tables.to(DEVICE) + storage = torch.randn(slots, 2, TPB, 3 * HD + 64, generator=g, device=DEVICE) + storage = storage.to(torch.bfloat16) + pool = storage[:, 1, :, : 3 * HD] + ape = (torch.randn(KPOOL, HD, generator=g, device=DEVICE) * 0.5).to(torch.bfloat16) + args = (pool, tables[lead:], positions, ape, TPB) + return args, dict(head_dim=HD, kpool=KPOOL, request_ids=request_ids), storage + + +def _key_scale(pool, tables, positions, request_ids): + """Per row: the ``(slot, row)`` of its pool's first member and the largest ``|k|`` among the + pool's members the row sees, the scale of the key it writes.""" + rows_per_table = positions.shape[1] if positions.dim() == 2 else 1 + pos = positions.reshape(-1) + order = torch.arange(pos.shape[0], device=DEVICE) + table = order // rows_per_table if request_ids is None else request_ids.long() + member = pos[:, None] // KPOOL * KPOOL + torch.arange(KPOOL, device=DEVICE) + slot = tables[table[:, None], (member // TPB).clamp(max=tables.shape[1] - 1)] + k = pool[slot, member % TPB, :HD].float().abs().amax(dim=-1) + k = torch.where(member <= pos[:, None], k, 0.0).amax(dim=-1) + return slot[:, 0], member[:, 0] % TPB, k + + +@pytest.mark.parametrize("cell", kpool_update.CELLS, ids=[c.why[:44] for c in kpool_update.CELLS]) +def test_kpool_update_matches_reference(cell, certifying): # noqa: F811 + """Only the pool-key columns of refreshed pools' first-member rows change, and each new + key is within ``2**-6 * K + 2**-7 * |key|`` of the reference's, ``K`` the largest ``|k|`` + the last refreshing row sees in the pool. + + Both sides round the same quantities to BF16 at the same points, but their exponentials + and sums differ in the last FP32 bits, which can move a probability or a product to the + neighbouring BF16 value. One step of a probability ``p`` moves a term by at most ``2**-7 * + p * |k|``, the product's rounding as much again: at most ``2**-6 * sum_m p_m |k_m| <= 2**-6 + * K`` over the members. The key's own rounding can take one more step, ``2**-7 * |key|``. + """ + certifying(kpool_update) + args, kwargs, storage = _args(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + before = storage.clone() + ref_storage = storage.clone() + assert kpool_update(*args, **kwargs) is None + kpool_update.reference(ref_storage[:, 1, :, : 3 * HD], *args[1:], **kwargs) + # Nothing but the pool-key columns of the layer's rows moves, on either side. + for s in (storage, ref_storage): + assert torch.equal(s[:, 0], before[:, 0]) + assert torch.equal(s[:, 1, :, : 2 * HD], before[:, 1, :, : 2 * HD]) + assert torch.equal(s[:, 1, :, 3 * HD :], before[:, 1, :, 3 * HD :]) + keys = (s[:, 1, :, 2 * HD : 3 * HD] for s in (storage, ref_storage, before)) + got, want, old = keys + slot, row, scale = _key_scale(before[:, 1, :, : 3 * HD], *args[1:3], kwargs["request_ids"]) + dest = slot * TPB + row + written = torch.zeros(got.shape[0] * TPB, dtype=torch.bool, device=DEVICE) + written[dest] = True + written = written.view(got.shape[:2]) + assert torch.equal(got[~written], old[~written]) + assert torch.equal(want[~written], old[~written]) + # The last row on a pool sees all the members the others do: the largest scale is its. + tol = torch.zeros(got.shape[0] * TPB, device=DEVICE) + tol = tol.scatter_reduce(0, dest, 2**-6 * scale, reduce="amax").view(got.shape[:2]) + err = (got.float() - want.float()).abs() + band = tol[..., None] + 2**-7 * torch.maximum(got.float().abs(), want.float().abs()) + assert bool((err <= band).all()), f"largest key error {float(err.max()):.3e}" + + +def test_guard_refuses_two_rows_on_one_pool(): + """Packed context rows run one program each: two rows on one pool race for its key.""" + args, kwargs, _ = _args(dict(context=((0, 64),), pages=8), seed=0) + pool, tables, _, ape, tpb = args + positions = torch.tensor([5, 6], dtype=torch.int64, device=DEVICE) + kwargs["request_ids"] = torch.zeros(2, dtype=torch.int32, device=DEVICE) + with pytest.raises(AssertionError, match="two rows refresh one pool"): + with validating(kpool_update): + kpool_update(pool, tables, positions, ape, tpb, **kwargs) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_mla_context_prologue.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_mla_context_prologue.py new file mode 100644 index 000000000000..24bd55d3b421 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_mla_context_prologue.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for the glm5_next fp8 tp4 entry mla_context_prologue.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.attention.mla_context_prologue import ( # noqa: E501 + mla_context_prologue, +) +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() != 100, reason="the glm5_next fp8 tp4 entries are certified on sm_100" +) + +DEVICE = "cuda" +TPB, PAGES = 32, 4100 +Q_LORA, KV_LORA, HD, INDEX_WIDTH = 1536, 512, 128, 288 +RMS_EPS, LN_EPS = 1e-5, 1e-6 + + +def _args(spec: dict, seed: int): + """The op's arguments for the packed context rows of ``(cached, new)`` requests, and the + storage behind both pool views. + + Every page a request reaches gets its own slot. The pools are coalesced views (the + latent pool the first of two FP8 subpages per slot, the indexer pool the second of three + BF16 subpages with padded rows) holding stale finite values. ``tables`` adds generation + tables after the context ones, as in a mixed step. + """ + g = torch.Generator(device=DEVICE).manual_seed(seed) + positions, request_ids, pages = [], [], [] + for i, (cached, new) in enumerate(spec["context"]): + positions.append(torch.arange(cached, cached + new)) + request_ids.append(torch.full((new,), i, dtype=torch.int32)) + pages.append(-(-(cached + new) // TPB)) + positions = torch.cat(positions).to(DEVICE) + request_ids = torch.cat(request_ids).to(DEVICE) + num_tables = max(spec.get("tables", 0), len(pages)) + tables = torch.zeros(num_tables, spec.get("pages", PAGES), dtype=torch.int64, device=DEVICE) + slots = 1 + sum(pages) + order = 1 + torch.randperm(slots - 1, generator=g, device=DEVICE) + start = 0 + for t, count in enumerate(pages): + tables[t, :count] = order[start : start + count] + start += count + rows = positions.shape[0] + qkv = torch.randn(rows, Q_LORA + KV_LORA, generator=g, device=DEVICE).to(torch.bfloat16) + idx = torch.randn(rows, INDEX_WIDTH, generator=g, device=DEVICE).to(torch.bfloat16) + weights = [ + (torch.randn(d, generator=g, device=DEVICE) * 0.2 + 1.0).to(torch.bfloat16) + for d in (Q_LORA, KV_LORA, HD) + ] + bias = (torch.randn(HD, generator=g, device=DEVICE) * 0.1).to(torch.bfloat16) + latent_storage = torch.randn(slots, 2, TPB, KV_LORA, generator=g, device=DEVICE) * 8 + latent_storage = latent_storage.to(torch.float8_e4m3fn) + index_storage = torch.randn(slots, 3, TPB, 3 * HD + 64, generator=g, device=DEVICE) + index_storage = index_storage.to(torch.bfloat16) + scale = torch.tensor(0.7, device=DEVICE) + args = [qkv, idx, *weights, bias, RMS_EPS, LN_EPS, positions, request_ids, tables] + return args, latent_storage, index_storage, scale + + +def _pools(latent_storage, index_storage): + return latent_storage[:, 0], index_storage[:, 1, :, : 3 * HD] + + +def _bytes(pool): + """A pool's rows as bytes: boolean indexing does not take every FP8 dtype.""" + return pool.view(torch.uint8) + + +@pytest.mark.parametrize( + "cell", mla_context_prologue.CELLS, ids=[c.why[:44] for c in mla_context_prologue.CELLS] +) +def test_mla_context_prologue_matches_reference(cell, certifying): # noqa: F811 + """The q_a rows, the latent rows and the indexer rows equal the reference's up to one + rounding step, nothing else moves. + + Both sides normalize in FP32 and round once, but the kernel's sums and reciprocal square + roots differ from torch's in the last FP32 bits, which can move a value to the + neighbouring BF16 one: at most one BF16 step, ``2**-7 |x|``, for the q_a rows and the + normalized indexer keys. An indexer key the LayerNorm bias cancels to near zero keeps the + FP32 error of the terms it cancelled, so its band has an absolute floor of ``2**-16`` (keys + are of order one; one such element of 552,064 differed by 3e-8 at 2.6e-6). The latent is + rounded to BF16 and then quantized: a BF16 step can move its E4M3 code by one step, at most + ``2**-3 |x|`` (``2**-9`` among subnormals). The gate columns are copied, bit for bit. + """ + certifying(mla_context_prologue) + args, latent_storage, index_storage, scale = _args( + cell.spec, zlib.crc32(cell.why.encode()) % 2**31 + ) + stale = (latent_storage.clone(), index_storage.clone()) + ref_storage = (latent_storage.clone(), index_storage.clone()) + q_resid = mla_context_prologue(*args, *_pools(latent_storage, index_storage), TPB, scale) + q_ref = mla_context_prologue.reference(*args, *_pools(*ref_storage), TPB, scale) + assert q_resid.shape == (args[0].shape[0], Q_LORA) and q_resid.dtype == torch.bfloat16 + assert bool(((q_resid.float() - q_ref.float()).abs() <= 2**-7 * q_ref.float().abs()).all()) + + positions, request_ids, tables = args[8], args[9], args[10] + slot = tables[request_ids.long(), positions // TPB] + row = positions % TPB + written = torch.zeros(latent_storage.shape[0], TPB, dtype=torch.bool, device=DEVICE) + written[slot, row] = True + # Untouched: the foreign subpages, the padding, the pool keys and every unwritten row. + assert torch.equal(_bytes(latent_storage[:, 1]), _bytes(stale[0][:, 1])) + for i in (0, 2): + assert torch.equal(index_storage[:, i], stale[1][:, i]) + assert torch.equal(index_storage[:, 1, :, 2 * HD :], stale[1][:, 1, :, 2 * HD :]) + for got, want, old in zip( + _pools(latent_storage, index_storage), _pools(*ref_storage), _pools(*stale) + ): + assert torch.equal(_bytes(got)[~written], _bytes(old)[~written]) + assert torch.equal(_bytes(want)[~written], _bytes(old)[~written]) + latent, latent_ref = _pools(latent_storage, index_storage)[0], _pools(*ref_storage)[0] + got = _bytes(latent)[written].view(torch.float8_e4m3fn).float() + want = _bytes(latent_ref)[written].view(torch.float8_e4m3fn).float() + assert bool(((got - want).abs() <= torch.clamp(2**-3 * want.abs(), min=2**-9)).all()) + index, index_ref = _pools(latent_storage, index_storage)[1], _pools(*ref_storage)[1] + got, want = index[written][:, :HD].float(), index_ref[written][:, :HD].float() + assert bool(((got - want).abs() <= torch.clamp(2**-7 * want.abs(), min=2**-16)).all()) + assert torch.equal(index[written][:, HD : 2 * HD], index_ref[written][:, HD : 2 * HD]) + + +def test_guard_refuses_a_position_written_twice(): + """Two rows on one cache position race on its latent and indexer rows.""" + args, latent_storage, index_storage, scale = _args(dict(context=((0, 8),), pages=4), seed=0) + args[8] = torch.tensor([0, 1, 2, 3, 4, 5, 6, 6], device=DEVICE) + with pytest.raises(AssertionError, match="written twice"): + with validating(mla_context_prologue): + mla_context_prologue(*args, *_pools(latent_storage, index_storage), TPB, scale) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_mla_decode_glue.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_mla_decode_glue.py new file mode 100644 index 000000000000..48ffe263eda6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_mla_decode_glue.py @@ -0,0 +1,125 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for the glm5_next fp8 tp4 MLA decode entries mla_prologue, kpool_select and +latent_values.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.attention.kpool_select import ( # noqa: E501 + kpool_select, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.attention.latent_values import ( # noqa: E501 + latent_values, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.attention.mla_prologue import ( # noqa: E501 + mla_prologue, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 100, reason="the glm5_next fp8 tp4 entries are certified on sm_100" +) + +_TPB, _SLOTS = 32, 512 + + +def _prologue_args(spec: dict, seed: int) -> list: + """``n`` requests of ``t`` consecutive rows each, at positions past the first two pages, with + the int64 GLM block tables the target passes.""" + g = torch.Generator(device="cuda").manual_seed(seed) + n, t = spec["n"], spec["t"] + + def rnd(*shape, scale=1.0): + return torch.randn(*shape, generator=g, device="cuda") * scale + + qkv = rnd(n * t, 2048).to(torch.bfloat16) + idx = rnd(n * t, 288).to(torch.bfloat16) + weights = [(rnd(d) * 0.2 + 1.0).to(torch.bfloat16) for d in (1536, 512, 128)] + bias = rnd(128, scale=0.1).to(torch.bfloat16) + positions = torch.randint(64, 600, (n, 1), generator=g, device="cuda") + torch.arange( + t, device="cuda" + ) + tables = torch.randperm(_SLOTS, generator=g, device="cuda")[: n * 20].view(n, 20) + latent_pool = torch.zeros(_SLOTS, _TPB, 512, device="cuda", dtype=torch.float8_e4m3fn) + index_pool = torch.zeros(_SLOTS, _TPB, 288, device="cuda", dtype=torch.bfloat16) + scale = torch.tensor(0.7, device="cuda") + return [ + qkv, + idx, + *weights, + bias, + 1e-5, + 1e-6, + positions, + tables, + latent_pool, + index_pool, + _TPB, + scale, + ] + + +@pytest.mark.parametrize("cell", mla_prologue.CELLS, ids=[c.why[:44] for c in mla_prologue.CELLS]) +def test_prologue_matches_reference(cell, certifying): # noqa: F811 + """The rows, the latent cache and the indexer cache equal the reference's (same FP32 math).""" + certifying(mla_prologue) + args = _prologue_args(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + mla_prologue.is_valid(*args) + ref_args = list(args) + ref_args[10], ref_args[11] = args[10].clone(), args[11].clone() + out = mla_prologue(*args) + ref = mla_prologue.reference(*ref_args) + torch.testing.assert_close( + out.float(), ref.float(), rtol=0.0, atol=2**-7 * ref.abs().max().item() + ) + assert torch.equal(args[10].view(torch.uint8) != 0, ref_args[10].view(torch.uint8) != 0) + diff = (args[10].float() - ref_args[10].float()).abs() + assert diff.max().item() <= 0.125 * ref_args[10].float().abs().max().item() + torch.testing.assert_close(args[11].float(), ref_args[11].float(), rtol=0.0, atol=2**-6) + + +@pytest.mark.parametrize("cell", latent_values.CELLS, ids=[c.why[:44] for c in latent_values.CELLS]) +def test_latent_values_matches_reference(cell, certifying): # noqa: F811 + """FP32 accumulation of BF16 products, one BF16 rounding; rows without a key are zero, NaN + latents included.""" + certifying(latent_values) + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + rows, heads, lat, v = (cell.spec[k] for k in ("r", "h", "l", "v")) + latent = torch.randn(rows, heads, lat, generator=g, device="cuda").to(torch.bfloat16) + w_v_t = (torch.randn(heads, lat, v, generator=g, device="cuda") * 0.05).to(torch.bfloat16) + keep = torch.ones(rows, dtype=torch.bool, device="cuda") + keep[-1] = False + latent[-1] = float("nan") + latent_values.is_valid(latent, w_v_t, keep) + out = latent_values(latent, w_v_t, keep) + ref = latent_values.reference(latent, w_v_t, keep) + assert torch.isfinite(out).all() + torch.testing.assert_close(out.float(), ref.float(), rtol=2**-7, atol=1e-3) + + +@pytest.mark.parametrize("cell", kpool_select.CELLS, ids=[c.why[:44] for c in kpool_select.CELLS]) +def test_kpool_select_matches_reference(cell, certifying): # noqa: F811 + """Packed rows, lengths and the nonempty mask equal top-k + expand + compaction of the pools + sorted by index; the last row sees no key.""" + certifying(kpool_select) + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + rows, tpb, kpool = spec["rows"], 32, 4 + blocks = spec["pools"] * kpool // tpb + tables = torch.randperm(2 * blocks, generator=g, device="cuda")[:blocks].repeat(rows, 1) + kv_lens = spec["kv"] - torch.arange(rows, device="cuda").flip(0) % 4 + kv_lens[-1] = 0 + scores = torch.randn(rows, spec["pools"], generator=g, device="cuda") + args = (scores, kv_lens, tables, tpb) + kwargs = dict(select_k=512, kpool=kpool, base_row=3, rows_per_slot=384, out_width=2051) + kpool_select.is_valid(*args, **kwargs) + out = kpool_select(*args, **kwargs) + ref = kpool_select.reference(*args, **kwargs) + for got, want in zip(out, ref): + assert torch.equal(got, want) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_sparse_mla_decode.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_sparse_mla_decode.py new file mode 100644 index 000000000000..f4889a0be40b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/attention/test_modeling_v2_sparse_mla_decode.py @@ -0,0 +1,147 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for the glm5_next fp8 tp4 entry sparse_mla_decode.""" + +import zlib + +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.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.attention.sparse_mla_decode import ( # noqa: E501 + WORKSPACE_BYTES, + counter_bytes, + sparse_mla_decode, +) +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() != 100, reason="the glm5_next fp8 tp4 entries are certified on sm_100" +) + +DEVICE = "cuda" +TPB, KPOOL, SELECT_K, WIDTH, DIM = 32, 4, 512, 2112, 512 +# A coalesced latent pool: the layer's 32 rows of a slot are the second of three subpages. +SLOTS, SUBPAGES, BASE_ROW = 512, 3, 32 +# The cache's dequantization scale and the model's softmax scale, qk_nope_head_dim ** -0.5. +KV_SCALE, SOFTMAX_SCALE = 0.75, 256**-0.5 + + +def _lengths(spec: dict) -> torch.Tensor: + """Each row's selection length: the members of its ``min(512, visible // 4)`` pools plus + its tail. Generation cells: request ``b``'s KV length after the step is ``kv - b % 4`` and + its rows see ``kv_b - T + 1 .. kv_b``; context cells: row ``r`` sees ``r + 1``.""" + if "context" in spec: + lo, hi = spec["rows"] + visible = torch.arange(lo + 1, hi + 1, device=DEVICE) + else: + kv = spec["kv"] - torch.arange(spec["requests"], device=DEVICE) % 4 + steps = torch.arange(spec["tokens"], device=DEVICE) + visible = (kv[:, None] - spec["tokens"] + 1 + steps).reshape(-1) + return ((visible // KPOOL).clamp(max=SELECT_K) * KPOOL + visible % KPOOL).to(torch.int32) + + +def _args(spec: dict, seed: int): + """The op's arguments: the rows of the coalesced FP8 pool (every row outside the layer's + subpage is NaN, except row 0, the dummy id), each row's distinct selected layer rows + packed in front of -1 padding, E4M3 queries, zeroed buffers, and the target's scales + (a host float for bmm1, the device dequantization scale for bmm2). ``empty_rows`` trailing + rows selected nothing: one dummy id 0, length 1.""" + g = torch.Generator(device=DEVICE).manual_seed(seed) + lengths = _lengths(spec) + rows, heads = lengths.shape[0], spec["heads"] + empty = spec.get("empty_rows", 0) + storage = torch.full((SLOTS, SUBPAGES, TPB, DIM), float("nan"), device=DEVICE) + storage[:, 1] = torch.randn(SLOTS, TPB, DIM, generator=g, device=DEVICE) + storage[0, 0] = torch.randn(TPB, DIM, generator=g, device=DEVICE) + kv_rows = storage.to(torch.float8_e4m3fn).view(-1, 1, DIM) + layer = BASE_ROW + torch.arange(SLOTS, device=DEVICE)[:, None] * SUBPAGES * TPB + layer = (layer + torch.arange(TPB, device=DEVICE)).reshape(-1) + pick = torch.rand(rows, layer.numel(), generator=g, device=DEVICE) + pick = pick.topk(int(lengths.max()), dim=1).indices + packed = torch.full((rows, WIDTH), -1, dtype=torch.int32, device=DEVICE) + packed[:, : pick.shape[1]] = layer[pick].to(torch.int32) + col = torch.arange(WIDTH, device=DEVICE) + packed = torch.where(col < lengths[:, None], packed, -1) + if empty: + packed[-empty:] = -1 + packed[-empty:, 0] = 0 + lengths[-empty:] = 1 + query = torch.randn(rows, heads, DIM, generator=g, device=DEVICE).to(torch.float8_e4m3fn) + sm_count = torch.cuda.get_device_properties(DEVICE).multi_processor_count + workspace = torch.zeros(WORKSPACE_BYTES, dtype=torch.uint8, device=DEVICE) + counter = torch.zeros(counter_bytes(rows, heads, sm_count), dtype=torch.uint8, device=DEVICE) + bmm2 = torch.tensor([KV_SCALE], dtype=torch.float32, device=DEVICE) + return query, kv_rows, packed, lengths, workspace, counter, KV_SCALE * SOFTMAX_SCALE, bmm2 + + +def _sigma(query, kv_rows, packed, lengths, ref, bmm1, bmm2): + """Standard deviation of each output's error from the E4M3 rounding of the probabilities. + + The kernel rounds ``p_j = exp(s_j - m) <= 1`` (``m`` a running maximum, never above the + row's) to E4M3 for the value product: at most half a step, ``2**-4 p_j`` above ``2**-6`` + and ``2**-10`` below, modelled as uniform (variance ``h**2 / 3``). In units of the + normalized ``pi_j = p_j / l`` that is ``h_j = max(2**-4 pi_j, 2**-10 / l)``. The output + ``o_d = sum_j pi_j v_jd`` then errs with variance ``sum_j h_j**2 v_jd**2 / 3``, plus + ``o_d**2 sum_j h_j**2 / 3`` should the normalizer use the rounded probabilities. + """ + rows, heads, dim = query.shape + latent = kv_rows.reshape(-1, dim).view(torch.uint8) + col = torch.arange(packed.shape[1], device=DEVICE) + var = torch.empty(rows, heads, dim, device=DEVICE) + for r0 in range(0, rows, 16): + part = slice(r0, min(r0 + 16, rows)) + selected = col[None, :] < lengths[part, None].long() + values = latent[packed[part].long().clamp(min=0)].view(kv_rows.dtype).float() + values = values.masked_fill(~selected[:, :, None], 0.0) + with true_fp32_matmul(): + logits = torch.einsum("rhd,rwd->rhw", query[part].float(), values) * bmm1 + logits = logits.masked_fill(~selected[:, None, :], float("-inf")) + p = (logits - logits.amax(dim=-1, keepdim=True)).exp() + norm = p.sum(dim=-1, keepdim=True) + h = torch.maximum(2**-4 * p / norm, 2**-10 / norm).masked_fill(~selected[:, None, :], 0) + with true_fp32_matmul(): + spread = torch.einsum("rhw,rwd->rhd", h**2, (values * bmm2) ** 2) + out = ref[part].float() + var[part] = (spread + out**2 * (h**2).sum(dim=-1, keepdim=True)) / 3 + return var.sqrt() + + +@pytest.mark.parametrize( + "cell", sparse_mla_decode.CELLS, ids=[c.why[:44] for c in sparse_mla_decode.CELLS] +) +def test_sparse_mla_decode_matches_reference(cell, certifying): # noqa: F811 + """Every output within ``8 sigma + 2**-7 |ref|`` of the FP32 reference: ``sigma`` from the + kernel's E4M3 probabilities (``_sigma``; the scores, the softmax and the accumulations are + FP32 on exact E4M3 operands), ``2**-7 |ref|`` for the two sides rounding to BF16 apart. + Across the cells' ~10**8 outputs a Gaussian tail past 8 sigma is below 1e-7. The relative + RMS error of the whole cell stays under ``2**-4`` (the rounding alone predicts about + ``2**-4 / sqrt(3)``). The kernel hands its multi-CTA counter back zeroed, and no NaN + outside the layer's rows is read.""" + certifying(sparse_mla_decode) + args = _args(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + out = sparse_mla_decode(*args) + ref = sparse_mla_decode.reference(*args) + query, kv_rows, packed, lengths, workspace, counter, bmm1, bmm2 = args + assert out.shape == ref.shape == query.shape and out.dtype == torch.bfloat16 + assert not bool(counter.any()), "the kernel leaves its counter at zero for the next launch" + assert bool(torch.isfinite(out).all()) + sigma = _sigma(query, kv_rows, packed, lengths, ref, bmm1, bmm2) + err = (out.float() - ref.float()).abs() + band = 8 * sigma + 2**-7 * ref.float().abs() + worst = float((err / band.clamp_min(1e-30)).max()) + assert bool((err <= band).all()), f"worst error {worst:.3f} of the band" + rel_rms = float(err.square().mean().sqrt() / ref.float().square().mean().sqrt()) + assert rel_rms < 2**-4, f"relative RMS error {rel_rms:.3e}" + + +def test_guard_refuses_a_short_workspace(): + """The kernel's split-KV scratch is the rest of the workspace: a short one is overrun.""" + args = list(_args(dict(heads=16, requests=1, tokens=4, kv=13312), seed=0)) + args[4] = torch.zeros(1 << 20, dtype=torch.uint8, device=DEVICE) + with pytest.raises(AssertionError, match="workspace"): + with validating(sparse_mla_decode): + sparse_mla_decode(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/comm/_mhc_ar_boundary_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/comm/_mhc_ar_boundary_op_matrix.py new file mode 100644 index 000000000000..511407670ece --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/comm/_mhc_ar_boundary_op_matrix.py @@ -0,0 +1,546 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix of this target's fused all-reduce + mHC boundary (`ops/comm/mhc_ar_boundary.py`). + +Same shape as the catalog's collective matrices (`tests/unittest/_torch/modeling_v2/comm/`): run plainly, this +script is its own launcher. It spawns one rank per device named in CUDA_VISIBLE_DEVICES (the op is compiled for +4) and kills the whole process group on the first failure or at the deadline, because a broken collective hangs +rather than raising. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _mhc_ar_boundary_op_matrix.py + +Every rank can regenerate every other rank's partial from the seed, so the reference is arithmetic rather than a +second collective. The IPC workspace's one host-side exchange of memory handles goes through a gloo group, as in +the catalog's allreduce matrix. + +Not a pytest module: the `check_*` bodies are one ordered sequence inside one job and read rank state that only +`_run_one_rank` binds. +""" + +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, List + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "the fused boundary requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 + +# Bound inside the rank body, never in the launcher: importing the entry pulls in tensorrt_llm, and the launcher +# must carry no CUDA or MPI state. +entry: Any = None +library_allreduce: Any = None +assert_within_ulp: Any = None +WORKSPACE: Any = None +RANK = 0 +WORLD = 1 +HOST_GROUP: Any = None +MAPPING: Any = None + +MULT, HIDDEN, MIX = 4, 4096, 24 +NORM_EPS, HC_EPS, SINKHORN_EPS, POST_MULT, LAYER_INPUT_EPS = 1e-6, 1e-6, 1e-6, 2.0, 1e-5 + +# The band `_compare` gates on, and how it was measured, is in its docstring. +BF16_ELEMENT_ULP, BF16_RMS_ULP, FP32_ATOL = 3.0, 0.5, 5e-3 +OUTPUTS = ("residual_out", "post_out", "comb_out", "layer_input") +# 1: rank 0 prints each call's deviations in the band's units; 2: prints them and gates nothing. +REPORT = os.environ.get("MHC_AR_MATRIX_REPORT", "0") + + +class _HostExchange: + """The one call `IpcMemory` makes on `Distributed`: a TP-group allgather of host objects.""" + + def tp_allgather(self, obj: Any) -> List[Any]: + gathered: List[Any] = [None] * WORLD + dist.all_gather_object(gathered, obj, group=HOST_GROUP) + return gathered + + +class _TunerExchange(_HostExchange): + """The collectives the autotuner's distributed tuning makes, over the host group.""" + + def tp_cp_allgather(self, obj: Any) -> List[Any]: + return self.tp_allgather(obj) + + def tp_barrier(self) -> None: + dist.barrier(group=HOST_GROUP) + + def tp_cp_broadcast(self, obj: Any, root: int = 0) -> Any: + box = [obj] + dist.broadcast_object_list(box, src=root, group=HOST_GROUP) + return box[0] + + +def _partial(rank: int, rows: int, seed: int) -> torch.Tensor: + """One rank's share of a block output, at the scale an o_proj / MoE partial has.""" + gen = torch.Generator(device="cuda").manual_seed(seed * 977 + rank + 1) + return (torch.randn(rows, HIDDEN, generator=gen, device="cuda") * 0.5).to(torch.bfloat16) + + +def _boundary_inputs(rows: int, seed: int) -> dict: + """The streams, the previous block's coefficients and the next block's mixer and norm (shared by the ranks).""" + gen = torch.Generator(device="cuda").manual_seed(seed * 31 + 7) + residual = torch.randn(rows, MULT, HIDDEN, generator=gen, device="cuda").to(torch.bfloat16) + post = torch.rand(rows, MULT, generator=gen, device="cuda") * POST_MULT + comb = torch.rand(rows, MULT, MULT, generator=gen, device="cuda") + comb = comb / comb.sum(-2, keepdim=True) + return dict( + residual=residual, + post=post, + comb=comb.contiguous(), + fn=torch.randn(MIX, MULT * HIDDEN, generator=gen, device="cuda") * 0.02, + scale=0.75 + 0.5 * torch.rand(3, generator=gen, device="cuda"), + base=torch.randn(MIX, generator=gen, device="cuda") * 0.1, + norm_weight=(1 + 0.1 * torch.randn(HIDDEN, generator=gen, device="cuda")).to( + torch.bfloat16 + ), + ) + + +def _scalars(cell) -> dict: + return dict( + sinkhorn_iters=cell.spec["sinkhorn_iters"], + norm_eps=NORM_EPS, + hc_eps=HC_EPS, + sinkhorn_eps=SINKHORN_EPS, + post_mult=POST_MULT, + layer_input_eps=LAYER_INPUT_EPS, + ) + + +def _report(out, ref) -> None: + """Print each output's deviation in the units `_compare` gates on.""" + # Imported here, not at module scope: _validating imports tensorrt_llm, which the launcher must not, and the + # rank body puts it on sys.path. + from _validating import ULP + + ulp = ULP[torch.bfloat16] + for name, o, r in zip(OUTPUTS, out, ref): + d = o.float() - r.float() + if r.dtype != torch.bfloat16: + print(f"[report] {name}: max abs {d.abs().max().item():.3g}", flush=True) + continue + flat_d, flat_r = d.reshape(-1, r.shape[-1]), r.float().reshape(-1, r.shape[-1]) + row_scale = flat_r.abs().amax(1, keepdim=True).clamp_min(1e-9) + element = (flat_d.abs() / row_scale).max().item() / ulp + rms = (flat_d.square().mean().sqrt() / flat_r.square().mean().sqrt()).item() / ulp + print(f"[report] {name}: {element:.2f} ulp per element, {rms:.3f} ulp RMS", flush=True) + + +def _compare(out, ref, layer_input_rms_ulp: float = BF16_RMS_ULP) -> None: + """The BF16 outputs within 3 ulp per element against the row scale and 0.5 ulp relative RMS, the FP32 + coefficients within 5e-3 absolute. + + Measured on B200 over the 15 calls of this file (4 ranks; the three cells x 3 seeds, 4 interleaved and 2 + graph-replayed calls): + + - residual_out: at most 1.74 ulp per element, 0.17 ulp RMS. The streams differ from the reference only by + the order of their FP32 sums before one BF16 rounding; a flip of one ulp of a row's largest elements + reads as 1 to 2 ulp of the row scale. + - layer_input: at most 2.02 ulp per element, 0.25 ulp RMS. A stream's flip carries into the collapsed + row on top of its own rounding, and `pre` comes out of the TF32 mixer. + - post_out and comb_out: at most 1.9e-3 and 1.4e-3 absolute, on values in [0, 2] and [0, 1]. The mixer + weights round to TF32 (the BF16 streams are exact in it). + + A reduction that drops a rank measures 148 to 175 ulp per element (120 to 130 ulp RMS) on the BF16 + outputs and 0.44 to 0.65 absolute on the coefficients; `check_the_gate_discriminates_a_wrong_reduction` + checks the band refuses one. + """ + if REPORT in ("1", "2") and RANK == 0: + _report(out, ref) + if REPORT == "2": + return + for name, o, r in zip(OUTPUTS, out, ref): + if r.dtype == torch.bfloat16: + rms_ulp = layer_input_rms_ulp if name == "layer_input" else BF16_RMS_ULP + assert_within_ulp(o, r, element_ulp=BF16_ELEMENT_ULP, rms_ulp=rms_ulp) + else: + torch.testing.assert_close( + o, r, rtol=0.0, atol=FP32_ATOL, msg=lambda m, n=name: f"{n}: {m}" + ) + + +def _run(rows: int, seed: int, cell, inputs: dict | None = None): + inputs = _boundary_inputs(rows, seed) if inputs is None else inputs + out = entry(_partial(RANK, rows, seed), WORKSPACE, RANK, **inputs, **_scalars(cell)) + ref = entry.reference( + [_partial(r, rows, seed) for r in range(WORLD)], WORKSPACE, RANK, **inputs, **_scalars(cell) + ) + return out, ref + + +def check_certified_cells() -> None: + """Every cell, three seeds, against the entry's reference and this file's band.""" + for cell in entry.CELLS: + for seed in range(3): + out, ref = _run(cell.spec["rows"], seed, cell) + _compare(out, ref) + + +def check_interleaves_with_the_library_allreduce() -> None: + """The op and the library's one-shot all-reduce share the workspace's Lamport rotation: alternate them.""" + cell = entry.CELLS[-1] + for seed in range(4): + out, ref = _run(cell.spec["rows"], seed, cell) + _compare(out, ref) + x = _partial(RANK, 8, 100 + seed) + summed = library_allreduce(x) + expect = torch.stack([_partial(r, 8, 100 + seed).float() for r in range(WORLD)]).sum(0) + assert_within_ulp( + summed, expect.to(torch.bfloat16), element_ulp=BF16_ELEMENT_ULP, rms_ulp=BF16_RMS_ULP + ) + + +def check_cuda_graph_replay() -> None: + """Captured once, replayed on new partials and streams written into the captured inputs.""" + cell = entry.CELLS[1] + rows = cell.spec["rows"] + inputs = _boundary_inputs(rows, 0) + partial = _partial(RANK, rows, 0) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + entry(partial, WORKSPACE, RANK, **inputs, **_scalars(cell)) + torch.cuda.current_stream().wait_stream(side) + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = entry(partial, WORKSPACE, RANK, **inputs, **_scalars(cell)) + for seed in (5, 6): + fresh = _boundary_inputs(rows, seed) + for name, tensor in fresh.items(): + inputs[name].copy_(tensor) + partial.copy_(_partial(RANK, rows, seed)) + dist.barrier() + graph.replay() + torch.cuda.synchronize() + ref = entry.reference( + [_partial(r, rows, seed) for r in range(WORLD)], + WORKSPACE, + RANK, + **inputs, + **_scalars(cell), + ) + _compare(out, ref) + + +def check_the_gate_discriminates_a_wrong_reduction() -> None: + """A reference that drops one rank's partial must fail the band.""" + cell = entry.CELLS[1] + rows, seed = cell.spec["rows"], 9 + inputs = _boundary_inputs(rows, seed) + out = entry(_partial(RANK, rows, seed), WORKSPACE, RANK, **inputs, **_scalars(cell)) + wrong = entry.reference( + [_partial(r, rows, seed) for r in range(WORLD - 1)], + WORKSPACE, + RANK, + **inputs, + **_scalars(cell), + ) + try: + _compare(out, wrong) + except AssertionError: + return + if REPORT != "2": + raise AssertionError("the band accepted a reduction that dropped a rank") + + +def _partial_with_zero_words(rank: int, rows: int, seed: int) -> torch.Tensor: + """``_partial`` with BF16 zero pairs in its first 192 columns: ``(+0.0, -0.0)`` pairs, whose 32-bit + word is the Lamport protocol's empty-slot sentinel ``0x80000000``, then ``(+0.0, +0.0)`` and + ``(-0.0, -0.0)`` pairs.""" + partial = _partial(rank, rows, seed) + partial[:, 0:64:2] = 0.0 + partial[:, 1:64:2] = -0.0 + partial[:, 64:128] = 0.0 + partial[:, 128:192] = -0.0 + return partial + + +def check_zero_words_do_not_stall_the_exchange() -> None: + """A partial whose 32-bit word equals the empty-slot sentinel is sent as +0.0, so the receivers' + polling does not wait on it: the call completes and matches the reference.""" + for cell in entry.CELLS: + rows, seed = cell.spec["rows"], 40 + inputs = _boundary_inputs(rows, seed) + partial = _partial_with_zero_words(RANK, rows, seed) + assert ( + partial[:, :64].view(torch.int32) + == torch.tensor(-(2**31), dtype=torch.int32, device="cuda") + ).all() + out = entry(partial, WORKSPACE, RANK, **inputs, **_scalars(cell)) + ref = entry.reference( + [_partial_with_zero_words(r, rows, seed) for r in range(WORLD)], + WORKSPACE, + RANK, + **inputs, + **_scalars(cell), + ) + _compare(out, ref) + + +def check_tuned_fused_hc_tactic_is_shared() -> None: + """A cold tuning pass of the decode boundary's ``mhc_fused_hc`` lands every TP rank on one tactic, so + the ranks' replicated streams (and the router logits computed from them) stay bitwise equal.""" + import zlib + + from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.mhc.mhc_fused_hc import ( # noqa: E501 + fused_hc_scratch, + fused_hc_tactic, + mhc_fused_hc, + ) + from tensorrt_llm._torch.autotuner import AutoTuner, autotune + + for name in ("FORCE_DETERMINISTIC", "FORCE_MHC_KERNEL_DETERMINISTIC"): + os.environ.pop(name, None) + rows, seed = 4, 50 + inputs = _boundary_inputs(rows, seed) + x = library_allreduce(_partial(RANK, rows, seed)) + args = ( + x, + inputs["residual"], + inputs["post"].unsqueeze(-1), + inputs["comb"], + inputs["fn"], + inputs["scale"], + inputs["base"], + inputs["norm_weight"], + entry.CELLS[0].spec["sinkhorn_iters"], + NORM_EPS, + HC_EPS, + SINKHORN_EPS, + POST_MULT, + LAYER_INPUT_EPS, + ) + tuner = AutoTuner.get() + tuner.clear_cache() + tuner.setup_distributed_state(MAPPING, _TunerExchange()) + try: + with autotune(): + tactic = fused_hc_tactic(*args) + tactics = _HostExchange().tp_allgather(tuple(tactic)) + assert all(t == tactics[0] for t in tactics), ( + f"the ranks tuned different tactics: {tactics}" + ) + out = mhc_fused_hc(*args, tactic, *fused_hc_scratch(rows, tactic, x.device)) + torch.cuda.synchronize() + digest = zlib.crc32( + b"".join(t.contiguous().view(torch.uint8).cpu().numpy().tobytes() for t in out) + ) + digests = _HostExchange().tp_allgather(digest) + assert len(set(digests)) == 1, f"the ranks' boundary outputs differ: {digests}" + finally: + tuner.clear_cache() + tuner.setup_distributed_state(MAPPING, None) + + +def check_matches_the_unfused_boundary() -> None: + """The decode target's wiring of the op against the path it replaces: the library one-shot all-reduce of + the partials, then the mHC site's ``fused_hc``, on a real ``Mhc`` site. + + The same band, except 1 ulp RMS on the layer input: the unfused path rounds the collapsed row to BF16 + before its RMSNorm and the op does not (see ``MhcArBoundary.reference``). Measured on B200 over the three + cells: the layer input within 1.95 ulp per element and 0.77 ulp RMS, the streams within 0.2 ulp, the + coefficients within 4e-5. + """ + from types import SimpleNamespace + + from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.layers.mhc import ( # noqa: E501 + Mhc, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.modeling import ( # noqa: E501 + DecodeTarget, + ) + + for cell in entry.CELLS: + rows, seed = cell.spec["rows"], 20 + inputs = _boundary_inputs(rows, seed) + dims = SimpleNamespace( + hc_mult=MULT, + hc_sinkhorn_iters=cell.spec["sinkhorn_iters"], + rms_eps=NORM_EPS, + hc_eps=HC_EPS, + ) + hc = Mhc(SimpleNamespace(fn=inputs["fn"], scale=inputs["scale"], base=inputs["base"]), dims) + target = SimpleNamespace( + comm=SimpleNamespace(workspace=WORKSPACE, tp_rank=RANK), + core=SimpleNamespace(dims=SimpleNamespace(rms_eps=LAYER_INPUT_EPS)), + ) + partial = _partial(RANK, rows, seed) + streams = (inputs["residual"], inputs["post"], inputs["comb"]) + fused = DecodeTarget._fused_boundary(target, partial, *streams, hc, inputs["norm_weight"]) + unfused = hc.fused_hc( + library_allreduce(partial), *streams, inputs["norm_weight"], LAYER_INPUT_EPS + ) + _compare( + fused, tuple(t.view_as(f) for t, f in zip(unfused, fused)), layer_input_rms_ulp=1.0 + ) + + +def check_is_valid_rejects_what_the_kernel_cannot_take() -> None: + cell = entry.CELLS[0] + inputs = _boundary_inputs(9, 0) + try: + entry.is_valid(_partial(RANK, 9, 0), WORKSPACE, RANK, **inputs, **_scalars(cell)) + except AssertionError: + return + raise AssertionError("is_valid accepted 9 rows") + + +CHECKS = ( + check_certified_cells, + check_interleaves_with_the_library_allreduce, + check_cuda_graph_replay, + check_the_gate_discriminates_a_wrong_reduction, + check_zero_words_do_not_stall_the_exchange, + check_tuned_fused_hc_tactic_is_shared, + check_matches_the_unfused_boundary, + check_is_valid_rejects_what_the_kernel_cannot_take, +) + + +def _run_one_rank() -> int: + """Body of one spawned rank: run every check; a failing rank exits nonzero and the launcher kills its + siblings.""" + global RANK, WORLD, HOST_GROUP, MAPPING, WORKSPACE, entry, library_allreduce, assert_within_ulp + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cuda:nccl,cpu:gloo", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + HOST_GROUP = dist.new_group(backend="gloo") + + from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm import allreduce as ar_entry + from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.comm import ( # noqa: E501 + mhc_ar_boundary as boundary_module, + ) + from tensorrt_llm._torch.distributed import communicator + from tensorrt_llm._utils import torch_pybind11_abi + from tensorrt_llm.bindings.internal.process_group import init_pg + from tensorrt_llm.functional import AllReduceFusionOp, AllReduceStrategy + from tensorrt_llm.mapping import Mapping + + assert world == boundary_module.WORLD, ( + f"the op is compiled for {boundary_module.WORLD} ranks, got {world}" + ) + local_group = dist.new_group(ranks=list(range(world)), backend="cuda:nccl,cpu:gloo") + init_pg(dist.group.WORLD, local_group, torch_pybind11_abi()) + communicator.Distributed.get = staticmethod(lambda mapping: _HostExchange()) + mapping = Mapping(world_size=world, rank=rank, gpus_per_node=world, tp_size=world) + MAPPING = mapping + + entry = boundary_module.mhc_ar_boundary + group_args = ar_entry.allreduce_group_args(mapping, strategy=int(AllReduceStrategy.ONESHOT)) + WORKSPACE = group_args["workspace"] + + def library_allreduce(x): + return ar_entry.allreduce( + x, + residual=None, + norm_weight=None, + scale=None, + bias=None, + op=int(AllReduceFusionOp.NONE), + eps=1e-6, + trigger_completion_at_end=False, + **group_args, + ) + + # Spawned worker, not a pytest module: reach the modeling_v2 test helpers directly. + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import assert_within_ulp as _assert_within_ulp + + assert_within_ulp = _assert_within_ulp + + 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 siblings; exit hard so the launcher kills them. + os._exit(13) + dist.barrier() + torch.cuda.synchronize() + print(f"[rank {RANK}] {len(CHECKS)} checks passed", flush=True) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker per visible device, rendezvous over env://, all killed on the first nonzero exit or at the + deadline.""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the 4 devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + 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", + ) + 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)), + 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/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/comm/test_modeling_v2_glm53_fp8_lamport_layout.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/comm/test_modeling_v2_glm53_fp8_lamport_layout.py new file mode 100644 index 000000000000..9a70a2389d5d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/comm/test_modeling_v2_glm53_fp8_lamport_layout.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The fused all-reduce boundary's copy of the library's Lamport protocol, pinned to the library. + +``csrc/mhcArBoundary.cu`` runs the one-shot Lamport all-reduce on the same workspace as the +library's, so its layout and rotation (namespace ``lamport``) must be the library's. The library +keeps them private: ``LamportComm`` in ``allReduceFusionKernels.cu`` (copied in +``MiniMaxReduceRMSKernel.cu``) reads them and ``allocate_allreduce_fusion_workspace`` builds them, +none of it in a header. This reads all of them, so changing either side alone fails here instead of +desynchronizing the two kernels at run time. A source check: it needs the repository's sources, not +a device. +""" + +import re +from pathlib import Path + +import pytest + +# tests/unittest/_torch/modeling_v2/models/glm5_next//comm -> repo root +_ROOT = Path(__file__).resolve().parents[8] +_TARGET_CSRC = ( + _ROOT + / "tensorrt_llm/_torch/_experimental/modeling_v2/models/glm5_next" + / "glm_5_3_flash_fp8__sm_10x__tp4/csrc/mhcArBoundary.cu" +) +_COMM_KERNELS = _ROOT / "cpp/tensorrt_llm/kernels/communicationKernels" +_LAMPORT_COMMS = ( + _COMM_KERNELS / "allReduceFusionKernels.cu", + _COMM_KERNELS / "MiniMaxReduceRMSKernel.cu", +) +_ALLOCATOR = _ROOT / "tensorrt_llm/_torch/distributed/allreduce_helper.py" + +# The target's constants (in terms of its group size kRanks); LamportComm's statements below set the same. +_CONSTANTS = { + "kBufferSlot": "2 * kRanks", + "kFlagSlot": "3 * kRanks", + "kLayoutSlot": "3 * kRanks + 1", + "kTableSize": "3 * kRanks + 2", + "kCounterWord": "0", + "kFlagWord": "2", + "kClearSizeWord": "0", + "kCommSizeWord": "1", + "kNumBuffers": "3", +} +# LamportComm's statements that set those values (whitespace-insensitive). +_LAMPORT_COMM_STATEMENTS = ( + # kFlagSlot, kCounterWord + "counter_ptr = &reinterpret_cast(workspace[NRanks * 3])[0];", + # kFlagSlot, kFlagWord + "flag_ptr = &reinterpret_cast(workspace[NRanks * 3])[2];", + # kLayoutSlot, kClearSizeWord + "clear_ptr = &reinterpret_cast(workspace[NRanks * 3 + 1])[0];", + # kLayoutSlot, kCommSizeWord + "auto comm_size = reinterpret_cast(workspace[NRanks * 3 + 1])[1];", + # dataBuffer + "int data_offset = flag_value % 3;", + # clearBuffer: the previous call's buffer + "int clear_offset = (flag_value + 2) % 3;", + # kBufferSlot + "data_bufs[r] = reinterpret_cast(workspace[2 * NRanks + r]) + data_offset * comm_size;", + "clear_buf = reinterpret_cast(workspace[2 * NRanks + rank]) + clear_offset * comm_size;", + # nextFlag, and the clear size and counter reset that go with it + "*flag_ptr = (flag_value + 1) % 3;", + "*clear_ptr = new_clear_size;", + "*counter_ptr = 0;", +) +# The one-shot kernel's clear size: the bytes it wrote, as the target writes it. +_ONE_SHOT_CLEAR_SIZE = ( + "comm.update(static_cast(params.size) * NRanks * sizeof(DType));", + "int clear_access = static_cast((comm.clear_size + sizeof(float4) - 1) / sizeof(float4));", +) +# The allocator's pointer table, flag words and layout words. +_ALLOCATOR_STATEMENTS = ( + "lamport_buffers = IpcMemory(mapping, 3 * lamport_buffers_size, is_p2p_supported)", + 'flag_buffer = torch.tensor([0, 0, 0], dtype=torch.int, device="cuda")', + 'layout_buffer = torch.tensor([0, lamport_buffers_size], dtype=torch.int64, device="cuda")', + "ipc_buffers.serialize() + ipc_barriers.serialize() + lamport_buffers.serialize()" + " + [flag_buffer.data_ptr()] + [layout_buffer.data_ptr()]", +) +# The target's rotation, in terms of its constants. +_TARGET_ROTATION = ( + "int dataBuffer(int flag) { return flag % kNumBuffers; }", + "int clearBuffer(int flag) { return (flag + kNumBuffers - 1) % kNumBuffers; }", + "int nextFlag(int flag) { return (flag + 1) % kNumBuffers; }", +) + +for _path in (_TARGET_CSRC, *_LAMPORT_COMMS, _ALLOCATOR): + if not _path.is_file(): + pytest.skip(f"the repository's sources are not here: no {_path}", allow_module_level=True) + + +def _squash(text: str) -> str: + """``text`` without whitespace, so formatting does not matter.""" + return re.sub(r"\s+", "", text) + + +def _lamport_namespace() -> str: + source = _TARGET_CSRC.read_text() + match = re.search(r"namespace lamport\s*\{(.*?)\}\s*// namespace lamport", source, re.S) + assert match, f"{_TARGET_CSRC.name} has no namespace lamport" + return match.group(1) + + +def _struct_lamport_comm(path: Path) -> str: + source = path.read_text() + match = re.search(r"struct LamportComm\s*\{(.*?)\n\};", source, re.S) + assert match, f"{path.name} has no struct LamportComm" + return match.group(1) + + +def test_target_constants(): + body = _lamport_namespace() + found = { + name: _squash(value) for name, value in re.findall(r"constexpr int (k\w+) = ([^;]+);", body) + } + assert found == {name: _squash(value) for name, value in _CONSTANTS.items()} + squashed = _squash(body) + for statement in _TARGET_ROTATION: + assert _squash(statement) in squashed, statement + + +@pytest.mark.parametrize("path", _LAMPORT_COMMS, ids=lambda p: p.name) +def test_library_lamport_comm(path): + squashed = _squash(_struct_lamport_comm(path)) + missing = [s for s in _LAMPORT_COMM_STATEMENTS if _squash(s) not in squashed] + assert not missing, ( + f"LamportComm in {path.name} changed; update namespace lamport in " + f"{_TARGET_CSRC.name} and this test together: {missing}" + ) + + +def test_library_one_shot_clear_size(): + squashed = _squash(_LAMPORT_COMMS[0].read_text()) + for statement in _ONE_SHOT_CLEAR_SIZE: + assert _squash(statement) in squashed, statement + # The target clears whole 16-byte stores, rounded up, and records the bytes it wrote. + target = _squash(_TARGET_CSRC.read_text()) + assert _squash("i < (clearSize + 15) / 16;") in target + assert ( + _squash( + "*clearPtr = static_cast(numRows) * kH * kRanks" + " * static_cast(sizeof(__nv_bfloat16));" + ) + in target + ) + + +def test_workspace_allocator(): + squashed = _squash(_ALLOCATOR.read_text()) + missing = [s for s in _ALLOCATOR_STATEMENTS if _squash(s) not in squashed] + assert not missing, ( + f"the all-reduce workspace layout in {_ALLOCATOR.name} changed; update namespace lamport " + f"in {_TARGET_CSRC.name} and this test together: {missing}" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/comm/test_modeling_v2_glm53_fp8_mhc_ar_boundary_checks.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/comm/test_modeling_v2_glm53_fp8_mhc_ar_boundary_checks.py new file mode 100644 index 000000000000..a8bcd085b1a3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/comm/test_modeling_v2_glm53_fp8_mhc_ar_boundary_checks.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The fused all-reduce boundary's binding refuses operands its kernel cannot read. + +Each case breaks one operand of an otherwise valid call; the binding must raise before the launch, +so one process and one device are enough (a call that launched would wait for the other ranks). +The collective itself is certified by ``_mhc_ar_boundary_op_matrix.py``. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4 import ( # noqa: E501 + native_ops, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.comm import ( # noqa: E501 + mhc_ar_boundary as boundary, +) +from tensorrt_llm._utils import get_sm_version + +pytestmark = pytest.mark.skipif( + get_sm_version() != 100, reason="the glm5_next fp8 tp4 entries are certified on sm_100" +) + +ROWS, MULT, HIDDEN = 4, boundary.MULT, boundary.HIDDEN +MIX = (2 + MULT) * MULT + + +def _call() -> dict: + """A call with every operand at the shape, dtype and layout the kernel is compiled for.""" + + def empty(*shape, dtype=torch.float32): + return torch.empty(shape, dtype=dtype, device="cuda") + + bf16 = torch.bfloat16 + return dict( + partial=empty(ROWS, HIDDEN, dtype=bf16), + workspace=torch.zeros(3 * boundary.WORLD + 2, dtype=torch.int64, device="cuda"), + rank=0, + residual=empty(ROWS, MULT, HIDDEN, dtype=bf16), + post=empty(ROWS, MULT), + comb=empty(ROWS, MULT, MULT), + fn=empty(MIX, MULT * HIDDEN), + scale=empty(3), + base=empty(MIX), + norm_weight=empty(HIDDEN, dtype=bf16), + sinkhorn_iters=20, + norm_eps=1e-6, + hc_eps=1e-6, + sinkhorn_eps=1e-6, + post_mult=2.0, + layer_input_eps=1e-5, + residual_out=empty(ROWS, MULT, HIDDEN, dtype=bf16), + post_out=empty(ROWS, MULT), + comb_out=empty(ROWS, MULT, MULT), + layer_input=empty(ROWS, HIDDEN, dtype=bf16), + pdl=False, + ) + + +def _misaligned(*shape, dtype): + """A contiguous tensor that starts 2 elements past a 16-byte boundary.""" + numel = 1 + for size in shape: + numel *= size + return torch.empty(numel + 8, dtype=dtype, device="cuda")[2 : 2 + numel].view(*shape) + + +_BAD = { + "partial": lambda: torch.empty(ROWS, HIDDEN, dtype=torch.float16, device="cuda"), + "workspace": lambda: torch.zeros(3 * boundary.WORLD + 1, dtype=torch.int64, device="cuda"), + "rank": lambda: boundary.WORLD, + "residual": lambda: torch.empty(ROWS, MULT + 1, HIDDEN, dtype=torch.bfloat16, device="cuda"), + "post": lambda: torch.empty(ROWS * MULT, device="cuda"), + "comb": lambda: torch.empty(ROWS, MULT, MULT, device="cuda").transpose(1, 2), + "fn": lambda: _misaligned(MIX, MULT * HIDDEN, dtype=torch.float32), + "scale": lambda: torch.empty(3, dtype=torch.bfloat16, device="cuda"), + "base": lambda: torch.empty(MIX - 1, device="cuda"), + "norm_weight": lambda: torch.empty(HIDDEN, device="cuda"), + "sinkhorn_iters": lambda: 0, + "residual_out": lambda: _misaligned(ROWS, MULT, HIDDEN, dtype=torch.bfloat16), + "post_out": lambda: torch.empty(ROWS, MULT, dtype=torch.bfloat16, device="cuda"), + "comb_out": lambda: torch.empty(ROWS, MULT * MULT, device="cuda"), + "layer_input": lambda: torch.empty(ROWS + 1, HIDDEN, dtype=torch.bfloat16, device="cuda"), +} + + +@pytest.mark.parametrize("name", list(_BAD)) +def test_rejects_an_operand_the_kernel_cannot_read(name): + call = _call() + call[name] = _BAD[name]() + with pytest.raises(RuntimeError, match=f"mhc_ar_boundary: {name} must"): + native_ops.mhc_ar_boundary(*call.values()) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/comm/test_modeling_v2_mhc_ar_boundary.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/comm/test_modeling_v2_mhc_ar_boundary.py new file mode 100644 index 000000000000..a081ca97fd0c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/comm/test_modeling_v2_mhc_ar_boundary.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the fused all-reduce + mHC boundary's 4-rank matrix. + +The matrix is ``_mhc_ar_boundary_op_matrix.py`` beside this file, its own 4-rank launcher; +``_rank_job`` (in the shared comm test directory) starts it in a fresh interpreter and turns +its exit code into an assertion. +""" + +from pathlib import Path + +import pytest +import torch + +__extra_import_path__ = ["../../../../comm"] +import _rank_job # noqa: E402 — needs the path declared above + + +# Spawns 4 rank processes over the visible devices. Under xdist several workers would fight for +# the same GPUs, so this must run alone, like the shared collective entries. +@pytest.mark.no_xdist +def test_mhc_ar_boundary_op_matrix() -> None: + devices = torch.cuda.device_count() if torch.cuda.is_available() else 0 + if devices < _rank_job.WORLD_SIZE: + pytest.skip( + f"the op runs on {_rank_job.WORLD_SIZE} ranks in one node; {devices} device(s) visible" + ) + claimed = [torch.cuda.get_device_capability(i) for i in range(_rank_job.WORLD_SIZE)] + if any(capability != (10, 0) for capability in claimed): + pytest.skip(f"the op is certified on sm_100; the claimed devices are {claimed}") + _rank_job.run("mhc_ar_boundary", directory=Path(__file__).resolve().parent) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/gemm/test_modeling_v2_glm53_fp8_dsv3_router_gemm.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/gemm/test_modeling_v2_glm53_fp8_dsv3_router_gemm.py new file mode 100644 index 000000000000..34cc526fd7af --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/gemm/test_modeling_v2_glm53_fp8_dsv3_router_gemm.py @@ -0,0 +1,139 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the dsv3_router_gemm operation of the glm5_next FP8 target. + +The entry owns the contract. This file turns a cell into the router's operands, gates the op +against the entry's float64 reference in fp32 ulp, and drives the bias guard. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.gemm.dsv3_router_gemm import ( # noqa: E501 + GLM_ROUTER, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.gemm.dsv3_router_gemm import ( # noqa: E501 + dsv3_router_gemm as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402,F401 — path above; fixture by name + ULP, + assert_within_ulp, + certifying, + validating, +) + +pytestmark = pytest.mark.skipif( + get_sm_version() != 100, reason="dsv3_router_gemm is certified on sm_100" +) + +DEV = "cuda" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """256 fp32 ulp of the row max element-wise, 64 fp32 ulp of relative RMS. + + The reference sums in float64, so the band measures only the op's own accumulation of + 4096 exact bf16 products. The small-M kernel adds in IEEE fp32: an fp32 simulation of its + order (32 sequential products per thread, butterfly, 4-warp sum) against a float64 sum + over 256 x 288 outputs of N(0, 1) bf16 operands reaches 1.2 ulp worst and 0.7 RMS. The + cuBLASLt path accumulates on the tensor cores, whose fp32 accumulation is coarser than + IEEE round-to-nearest: on B200 it reached 74 ulp worst (13,312 rows) and 17.7 ulp RMS + (256 rows) -- more than a fully sequential IEEE sum, which the same simulation puts at + 17.6 / 5.2. The band leaves 3.5x over that and still rejects a bf16-rounded output (one + bf16 ulp is 65536 fp32 ulp) or a TF32 product (thousands of ulp RMS) by orders of + magnitude. + """ + assert out.dtype == ref.dtype == torch.float32, (out.dtype, ref.dtype) + assert_within_ulp(out, ref, element_ulp=256.0, rms_ulp=64.0, ulp=ULP[torch.float32]) + + +def _build(spec: dict, seed: int) -> tuple[torch.Tensor, torch.Tensor]: + """BF16 hidden rows and the router weight's `.t()`: the weight is built `[N, K]` and + transposed, which is what makes `mat_b` column-major without a copy.""" + g = torch.Generator(device=DEV).manual_seed(seed) + hidden = torch.randn(spec["rows"], spec["hidden"], generator=g, device=DEV) + weight = torch.randn(spec["experts"], spec["hidden"], generator=g, device=DEV) + return hidden.to(torch.bfloat16), weight.to(torch.bfloat16).t() + + +# The shapes the target calls this operation with, on top of the operation's own `CELLS`. +# +# GLM-5.3-Flash at TP4: the router is replicated, so every rank routes its rows over all 288 +# experts from the full hidden width. Generation steps are MTP-3 verifies of 4 rows per +# request (up to 64 requests) and the drafter's single rows; a 13,312-row prefill runs +# sequence parallel, each rank routing its own quarter (3,328 rows). +CELLS: tuple[Cell, ...] = op.CELLS + ( + Cell(why="one row: the MTP drafter's decode step", spec=dict(rows=1, **GLM_ROUTER)), + Cell( + why="MTP-3 verify of one request: 4 rows on the small-M kernel", + spec=dict(rows=4, **GLM_ROUTER), + ), + Cell( + why="MTP-3 verify of two requests: 8 rows on the small-M kernel", + spec=dict(rows=8, **GLM_ROUTER), + ), + Cell( + why="MTP-3 verify of 8 requests: 32 rows on cuBLASLt", + spec=dict(rows=32, **GLM_ROUTER), + ), + Cell( + why="MTP-3 verify of 64 requests: 256 rows on cuBLASLt", + spec=dict(rows=256, **GLM_ROUTER), + ), + Cell( + why="sequence-parallel prefill: one rank's 3,328 rows of a 13,312-row step", + spec=dict(rows=3328, **GLM_ROUTER), + ), + Cell( + why="a 13,312-row prefill step routed in one call", + spec=dict(rows=13312, **GLM_ROUTER), + ), +) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture + certifying(op) + mat_a, mat_b = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + out = op(mat_a, mat_b, None, torch.float32) + assert out.shape == (cell.spec["rows"], cell.spec["experts"]) and out.is_contiguous() + _compare(out, op.reference(mat_a, mat_b, None, torch.float32)) + + +def test_padded_operands_are_repacked() -> None: + """A hidden row read out of a wider buffer, and a router weight with padded rows, are + copied into the dense layout before either kernel reads them -- no guard is needed.""" + g = torch.Generator(device=DEV).manual_seed(3) + wide = torch.randn(8, 4096 + 64, generator=g, device=DEV).to(torch.bfloat16) + padded = torch.randn(288, 4096 + 64, generator=g, device=DEV).to(torch.bfloat16) + mat_a, mat_b = wide[:, :4096], padded[:, :4096].t() + assert mat_a.stride(0) != 4096 and mat_b.stride(1) != 4096 + with validating(op): + out = op(mat_a, mat_b, None, torch.float32) + _compare(out, op.reference(mat_a, mat_b, None, torch.float32)) + + +def _ok() -> tuple[torch.Tensor, torch.Tensor]: + return _build(dict(rows=4, hidden=256, experts=32), seed=5) + + +def test_guard_refuses_a_bias_in_the_wrong_dtype() -> None: + mat_a, mat_b = _ok() + bias = torch.randn(32, dtype=torch.bfloat16, device=DEV) + with pytest.raises(AssertionError, match="bias dtype"): + with validating(op): + op(mat_a, mat_b, bias, torch.float32) + + +def test_guard_refuses_a_bias_of_the_wrong_shape() -> None: + mat_a, mat_b = _ok() + bias = torch.randn(64, dtype=torch.float32, device=DEV)[::2] + with pytest.raises(AssertionError, match="contiguous \\[N\\] tensor"): + with validating(op): + op(mat_a, mat_b, bias, torch.float32) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/gemm/test_modeling_v2_glm53_fp8_fp8_blockwise_gemm.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/gemm/test_modeling_v2_glm53_fp8_fp8_blockwise_gemm.py new file mode 100644 index 000000000000..39663cc8dfe6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/gemm/test_modeling_v2_glm53_fp8_fp8_blockwise_gemm.py @@ -0,0 +1,213 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the fp8_blockwise_gemm operation of the glm5_next FP8 target. + +The entry owns the contract. This file turns a cell into block-FP8 operands (drawn directly as +e4m3 codes and fp32 scales, so both sides read bit-identical values), runs the block-FP8 +`Linear` composition once end to end from bf16 rows, and drives the guard. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.gemm.fp8_blockwise_gemm import ( # noqa: E501 + BLOCK, + GLM_DENSE_DOWN, + GLM_DENSE_GATE_UP, + GLM_SHARED_DOWN, + GLM_SHARED_GATE_UP, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.gemm.fp8_blockwise_gemm import ( # noqa: E501 + fp8_blockwise_gemm as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.quantization.fp8_quantize_1x128 import ( # noqa: E501 + fp8_quantize_1x128, +) +from tensorrt_llm._torch.autotuner import AutoTuner, autotune +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402,F401 — path above; fixture by name + assert_within_ulp, + certifying, + validating, +) + +pytestmark = pytest.mark.skipif( + get_sm_version() != 100, reason="fp8_blockwise_gemm is certified on sm_100" +) + +DEV = "cuda" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """4 bf16 ulp of the row max element-wise, 2 ulp of relative RMS. + + Kernel and reference read bit-identical e4m3 codes and fp32 scales; the products of codes + are exact in fp32, and the two differ only in where the scales are applied (per 128-wide K + block in the kernel, before the sum in the reference) and in the order of the fp32 sums -- + a few fp32 ulp of the row scale, far below one bf16 ulp. What remains is the single bf16 + rounding on each side landing on different neighbours, one ulp of an element at most. A + misread scale layout (a transposed activation scale, a 64-column block offset) or a wrong + K block is off by a whole block's contribution, many ulp RMS. + """ + assert out.dtype == ref.dtype == torch.bfloat16, (out.dtype, ref.dtype) + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0) + + +def _blocks(n: int) -> int: + return (n + BLOCK - 1) // BLOCK + + +def _build(spec: dict, seed: int) -> tuple[torch.Tensor, ...]: + """Block-FP8 operands at O(1) output scale, every value exact in its storage type. + + Activation codes ~ 2 N(0, 1) in e4m3 with scales in [0.5, 1.5); weight codes ~ N(0, 1) in + e4m3 with scales in [0.5, 1.5) / sqrt(K), so the outputs are O(1) at every K. + """ + g = torch.Generator(device=DEV).manual_seed(seed) + m, n, k = spec["m"], spec["n"], spec["k"] + x = (2.0 * torch.randn(m, k, generator=g, device=DEV)).to(torch.float8_e4m3fn) + x_scale = 0.5 + torch.rand(_blocks(k), m, generator=g, device=DEV) + w = torch.randn(n, k, generator=g, device=DEV).to(torch.float8_e4m3fn) + w_scale = (0.5 + torch.rand(_blocks(n), _blocks(k), generator=g, device=DEV)) / k**0.5 + return x, w, x_scale, w_scale + + +# The shapes the target calls this operation with, on top of the operation's own `CELLS`. +# +# GLM-5.3-Flash at TP4 runs four block-FP8 FFN projections -- the three dense layers' MLP and +# every MoE layer's shared expert, each a gate_up and a down -- and, in the sequence-parallel +# MLA prefill, q_b and o over all 64 heads for a rank's own rows. Generation steps are MTP-3 +# verifies of 4 rows per request, up to 64 requests; a 13,312-row prefill runs the FFN GEMMs +# over every gathered row and the MLA ones over a rank's quarter (3,328 rows). +_FFN = ( + ("dense gate_up", GLM_DENSE_GATE_UP), + ("dense down", GLM_DENSE_DOWN), + ("shared gate_up", GLM_SHARED_GATE_UP), + ("shared down", GLM_SHARED_DOWN), +) +_ROWS = ( + (4, "MTP-3 verify of one request"), + (256, "MTP-3 verify of 64 requests"), + (13312, "a 13,312-row prefill step"), +) + +CELLS: tuple[Cell, ...] = op.CELLS + ( + *( + Cell( + why=f"{name} {shape['n']}x{shape['k']}, {m} rows: {what}", + spec=dict(m=m, **shape), + ) + for name, shape in _FFN + for m, what in _ROWS + ), + Cell( + why="sequence-parallel MLA prefill q_b over 64 heads (16384 x 1536), 3,328 rows", + spec=dict(m=3328, n=16384, k=1536), + ), + Cell( + why="sequence-parallel MLA prefill o over 64 heads (4096 x 16384), 3,328 rows", + spec=dict(m=3328, n=4096, k=16384), + ), +) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture + certifying(op) + x, w, x_scale, w_scale = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + out = op(x, w, x_scale, w_scale) + assert out.shape == (cell.spec["m"], cell.spec["n"]) and out.is_contiguous() + _compare(out, op.reference(x, w, x_scale, w_scale)) + + +def test_block_fp8_linear_from_bf16_rows() -> None: + """The composition a block-FP8 `Linear` runs: the op on `fp8_quantize_1x128`'s output, + passed straight through. Against the entry's reference on that same quantization it holds + the certified band, which is what shows the quantizer's scale layout is the one the GEMM + reads.""" + g = torch.Generator(device=DEV).manual_seed(23) + rows = torch.randn(64, 4096, generator=g, device=DEV).to(torch.bfloat16) + x, x_scale = fp8_quantize_1x128(rows) + _, w, _, w_scale = _build(dict(m=1, **GLM_SHARED_GATE_UP), seed=29) + with validating(op): + out = op(x, w, x_scale, w_scale) + _compare(out, op.reference(x, w, x_scale, w_scale)) + + +def test_tuned_tactic_under_cuda_graph_capture() -> None: + """A decode step's FFN input as the target runs it: `fp8_quantize_1x128`, then the op at the + tactic the autotuner picked in warmup, captured in a CUDA graph and replayed on new rows. + Tuning profiles only the call's own row bucket. The replayed quantization equals an eager + one bit for bit, and the product holds the certified band against the reference on it.""" + shape = GLM_SHARED_GATE_UP + g = torch.Generator(device=DEV).manual_seed(47) + rows = torch.randn(32, shape["k"], generator=g, device=DEV).to(torch.bfloat16) + _, w, _, w_scale = _build(dict(m=1, **shape), seed=53) + tuner = AutoTuner.get() + tuner.clear_cache() + try: + x, x_scale = fp8_quantize_1x128(rows) + with autotune(skip_dynamic_tuning_buckets=True): + op(x, w, x_scale, w_scale) + assert any( + key[0] == "trtllm::cute_dsl_fp8_gemm_blackwell::gemm" + for key in tuner.profiling_cache.cache + ), "the warmup tuned nothing" + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + warm_x, warm_scale = fp8_quantize_1x128(rows) + op(warm_x, w, warm_scale, w_scale) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + x, x_scale = fp8_quantize_1x128(rows) + out = op(x, w, x_scale, w_scale) + for seed in (59, 61): + fresh = torch.Generator(device=DEV).manual_seed(seed) + rows.copy_(torch.randn(rows.shape, generator=fresh, device=DEV).to(torch.bfloat16)) + graph.replay() + torch.cuda.synchronize() + eager_x, eager_scale = fp8_quantize_1x128(rows) + assert torch.equal(x.view(torch.uint8), eager_x.view(torch.uint8)) + assert torch.equal(x_scale, eager_scale) + _compare(out, op.reference(x, w, x_scale, w_scale)) + finally: + tuner.clear_cache() + + +def test_guard_refuses_a_transposed_activation_scale() -> None: + """A token-major `[M, K/128]` scale where the shapes differ is refused before the kernel + misreads it; at M == K/128 nothing can tell, which is why the entry says so.""" + x, w, x_scale, w_scale = _build(dict(m=8, **GLM_SHARED_GATE_UP), seed=31) + with pytest.raises(AssertionError, match="K-block-major"): + with validating(op): + op(x, w, x_scale.t().contiguous(), w_scale) + + +def test_guard_refuses_a_strided_scale_slice() -> None: + """A row slice of the activation needs a contiguous copy of its scale columns.""" + x, w, x_scale, w_scale = _build(dict(m=16, **GLM_SHARED_GATE_UP), seed=37) + with pytest.raises(AssertionError, match="input_scale must be contiguous"): + with validating(op): + op(x[4:12], w, x_scale[:, 4:12], w_scale) + + +def test_guard_refuses_bf16_scales() -> None: + x, w, x_scale, w_scale = _build(dict(m=8, **GLM_SHARED_GATE_UP), seed=41) + with pytest.raises(AssertionError, match="both scales must be fp32"): + with validating(op): + op(x, w, x_scale, w_scale.to(torch.bfloat16)) + + +def test_guard_refuses_mismatched_k() -> None: + x, w, x_scale, w_scale = _build(dict(m=8, n=1024, k=4096), seed=43) + narrow = x[:, :3072].contiguous() + with pytest.raises(AssertionError, match="takes K from"): + with validating(op): + op(narrow, w, x_scale[:24].contiguous(), w_scale) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/gemm/test_modeling_v2_low_m_bf16_gemm.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/gemm/test_modeling_v2_low_m_bf16_gemm.py new file mode 100644 index 000000000000..14960d1f23e3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/gemm/test_modeling_v2_low_m_bf16_gemm.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm5_next fp8 tp4 low_m_bf16_gemm entry. + +Every row count a table entry runs compiles its own CuTe DSL kernel, seconds each. The default +run drives the row counts at the tactic boundaries of each shape; ``LOW_M_BF16_GEMM_TEST_ALL_M=1`` +drives every row count from 1 to 32. +""" + +import os +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.gemm.low_m_bf16_gemm import ( # noqa: E501 + _CHOICES, + low_m_choice, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.gemm.low_m_bf16_gemm import ( # noqa: E501 + low_m_bf16_gemm as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 100, reason="the glm5_next fp8 tp4 entries are certified on sm_100" +) + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Kernel and reference multiply the same BF16 operands exactly in FP32 and differ only in + the order of the FP32 sums (split-K slices, the direct kernel's per-thread partials, cuBLAS), + about 1e-6 of a row's scale: at most rare one-ulp flips of the single BF16 rounding. 4 ulp + element-wise and 2 ulp RMS admit those; a lost K slice or wrong tile is off by hundreds.""" + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0) + + +def _build(m: int, n: int, k: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + 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) + return x, torch.nn.Parameter(weight, requires_grad=False) + + +@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) + s = cell.spec + x, weight = _build(s["m"], s["n"], s["k"], seed=zlib.crc32(cell.why.encode()) % 2**31) + _compare(op(x, weight), op.reference(x, weight.detach())) + + +#: The largest row count a table entry serves. +_MAX_M = 32 + +_ALL_M = os.environ.get("LOW_M_BF16_GEMM_TEST_ALL_M", "0") == "1" + + +def _boundary_row_counts(n: int, k: int) -> list[int]: + """The first and last row count of every run of row counts that take the same choice + (kernel and tactic, or cuBLAS), up to ``_MAX_M``: each tactic at its smallest and largest + M, both sides of every switch, and a direct tactic's row tile skipped where it does not + divide M.""" + choices = [low_m_choice(m, n, k) for m in range(1, _MAX_M + 1)] + ends = {1, _MAX_M} + for m in range(2, _MAX_M + 1): + if choices[m - 1] != choices[m - 2]: + ends.update((m - 1, m)) + return sorted(ends) + + +@pytest.mark.parametrize("shape", sorted(_CHOICES), ids=lambda s: f"{s[0]}x{s[1]}") +def test_row_counts(shape: tuple[int, int]) -> None: + """Every row count driven picks a tactic its kernel accepts (direct tactics need their row + tile to divide M) and computes the product: the tactic boundaries, or with + ``LOW_M_BF16_GEMM_TEST_ALL_M=1`` every row count of the table's range.""" + n, k = shape + row_counts = range(1, _MAX_M + 1) if _ALL_M else _boundary_row_counts(n, k) + for m in row_counts: + x, weight = _build(m, n, k, seed=m) + choice = low_m_choice(m, n, k) + if choice is not None and choice[0] == "direct": + rows_per_block = choice[1][2] + assert rows_per_block <= m and m % rows_per_block == 0 + _compare(op(x, weight), op.reference(x, weight.detach())) + + +def test_cuda_graph_capture_and_replay() -> None: + """The decode step is captured: the op must replay on new contents of the same input.""" + s = op.CELLS[1].spec + x, weight = _build(s["m"], s["n"], s["k"], seed=31) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(x, weight) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(x, weight) + fresh, _ = _build(s["m"], s["n"], s["k"], seed=32) + x.copy_(fresh) + graph.replay() + torch.cuda.synchronize() + _compare(out, op.reference(x, weight.detach())) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/mhc/_glm53_mhc_inputs.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/mhc/_glm53_mhc_inputs.py new file mode 100644 index 000000000000..25e48f4404a9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/mhc/_glm53_mhc_inputs.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Inputs of the glm5_next fp8 tp4 mHC op tests, at GLM-5.3-Flash's scale. + +The checkpoint stores each sublayer's mixer ``fn [24, 16384]`` as BF16 values (standard deviation +0.02 to 0.03 per layer) and its ``scale`` (0.04 to 0.6) and ``base`` in FP32; the target holds +``fn`` as FP32. The ``fn`` drawn here is BF16-representable like the checkpoint's, which TF32 holds +exactly, so the tensor-core backends differ from the references in summation order only, as they +do on the real weights. ``base`` is drawn wider than the checkpoint's (whose combination logits sit +near a diagonal), so that every combination entry and every stream carries weight and a transposed +or misplaced coefficient shows. +""" + +import torch + +MULT, HIDDEN, MIX = 4, 4096, 24 + +#: The op arguments after the tensors, in call order. +SCALAR_NAMES = ( + "sinkhorn_iters", + "norm_eps", + "hc_eps", + "sinkhorn_eps", + "post_mult", + "layer_input_eps", +) + + +def generator(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def scalars(spec) -> list: + """A cell's mHC scalars in call order.""" + return [spec[name] for name in SCALAR_NAMES] + + +def streams(rows: int, gen: torch.Generator, std: float = 1.0) -> torch.Tensor: + """``[rows, 4, 4096]`` BF16 residual streams.""" + return (torch.randn(rows, MULT, HIDDEN, generator=gen, device="cuda") * std).to(torch.bfloat16) + + +def block_output(rows: int, gen: torch.Generator, std: float = 1.0) -> torch.Tensor: + """``[rows, 4096]`` BF16: a sublayer's reduced output.""" + return (torch.randn(rows, HIDDEN, generator=gen, device="cuda") * std).to(torch.bfloat16) + + +def coefficients(rows: int, gen: torch.Generator) -> tuple[torch.Tensor, torch.Tensor]: + """A sublayer's ``post [rows, 4]`` in ``[0, 2]`` and column-normalized ``comb [rows, 4, 4]``.""" + post = torch.rand(rows, MULT, generator=gen, device="cuda") * 2 + comb = torch.rand(rows, MULT, MULT, generator=gen, device="cuda") + 0.05 + return post, comb / comb.sum(-2, keepdim=True) + + +def mixer(gen: torch.Generator) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """A sublayer's ``(fn [24, 16384], scale [3], base [24])``, FP32.""" + fn = torch.randn(MIX, MULT * HIDDEN, generator=gen, device="cuda") * 0.03 + scale = torch.tensor([0.1, 0.08, 0.3], device="cuda") + base = torch.randn(MIX, generator=gen, device="cuda") * 2 + return fn.to(torch.bfloat16).float(), scale, base + + +def norm_weight(gen: torch.Generator) -> torch.Tensor: + """A sublayer's input RMSNorm weight ``[4096]``, BF16.""" + return (1 + 0.1 * torch.randn(HIDDEN, generator=gen, device="cuda")).to(torch.bfloat16) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/mhc/test_modeling_v2_mhc_big_fuse.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/mhc/test_modeling_v2_mhc_big_fuse.py new file mode 100644 index 000000000000..524b90256a48 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/mhc/test_modeling_v2_mhc_big_fuse.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm5_next fp8 tp4 mhc_big_fuse entry. + +The inputs are the reference mixer product of the cell's streams, so the kernel and the reference +start from the same FP32 values and only this op's arithmetic is compared. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.mhc.mhc_big_fuse import ( # noqa: E501 + BLOCK_SIZE_OPTIONS, + big_fuse_block_size, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.mhc.mhc_big_fuse import ( # noqa: E501 + mhc_big_fuse as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.mhc.mhc_gemm_sqrsum import ( # noqa: E501 + mhc_gemm_sqrsum, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = [".", "../../../.."] +from _glm53_mhc_inputs import generator, mixer, norm_weight, scalars, streams # noqa: E402 +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 100, reason="the glm5_next fp8 tp4 entries are certified on sm_100" +) + +#: Absolute band of the FP32 coefficients; see `_compare`. +COEFF_ATOL = 1e-4 + + +def _compare(out, ref, fused_norm: bool) -> None: + """``post`` and ``comb`` within 1e-4 absolute; the layer input within 3 ulp per element of + its row scale (4 with the fused norm) and 0.5 ulp relative RMS. + + The coefficients come from the same mixer product on both sides; the kernel's fast-math + ``exp``, ``rsqrt`` and divisions are each a few FP32 ulp off, and the Sinkhorn tail repeats + its divisions 39 times, so the two differ by about 1e-6 -- 1e-4 is two orders above that, on + values in ``[0, 2]`` and ``[0, 1]``, and a misplaced scale, base or multiplier moves them by + 0.01 to 1. The layer input is rounded to BF16 at the same places on both sides (the collapse, + and with the norm the normalized product), so it differs only where those FP32 differences + flip a rounding: one BF16 step, at most 2 ulp of the row scale, or two steps with the norm + when both roundings flip. + """ + post, comb, layer_input = out + ref_post, ref_comb, ref_layer_input = ref + assert post.shape == ref_post.shape and comb.shape == ref_comb.shape + torch.testing.assert_close(post, ref_post, rtol=0.0, atol=COEFF_ATOL) + torch.testing.assert_close(comb, ref_comb, rtol=0.0, atol=COEFF_ATOL) + element_ulp = 4.0 if fused_norm else 3.0 + assert_within_ulp(layer_input, ref_layer_input, element_ulp=element_ulp, rms_ulp=0.5) + + +def _matches(out, ref, fused_norm: bool) -> bool: + try: + _compare(out, ref, fused_norm) + except AssertionError: + return False + return True + + +def _inputs(rows: int, norm: bool, seed: int, std: float = 1.0): + gen = generator(seed) + residual = streams(rows, gen, std) + fn, scale, base = mixer(gen) + mixes, sqrsum = mhc_gemm_sqrsum.reference(residual.view(rows, -1), fn, 0) + return [mixes, sqrsum, residual, scale, base, norm_weight(gen) if norm else None] + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + """Every block size, against the reference; the coefficients do not depend on it.""" + certifying(op) + spec = cell.spec + args = _inputs(spec["rows"], spec["norm"], seed=zlib.crc32(cell.why.encode()) % 2**31) + args += scalars(spec) + ref = op.reference(*args, big_fuse_block_size(spec["rows"])) + first = None + for block_size in BLOCK_SIZE_OPTIONS: + out = op(*args, block_size) + _compare(out, ref, spec["norm"]) + if first is None: + first = out + continue + assert torch.equal(out[0], first[0]) and torch.equal(out[1], first[1]), ( + f"block_size={block_size} changed the coefficients" + ) + if not spec["norm"]: + assert torch.equal(out[2], first[2]), f"block_size={block_size} changed the collapse" + + +def test_scalars_reach_their_kernel_parameters() -> None: + """Distinct, exaggerated scalars on small streams, where every epsilon moves the outputs. + + GLM-5.3-Flash's two epsilon pairs are equal (``norm_eps`` and ``layer_input_eps`` 1e-5, + ``hc_eps`` and ``sinkhorn_eps`` 1e-6) and its epsilons are too small to move unit-scale + streams, so the cells cannot tell a swapped pair from a right one. Here the streams' mean + square is 2.5e-3, the size of ``norm_eps``; the references with either pair swapped must fail + the band, and the op must pass it. + """ + args = _inputs(64, norm=True, seed=11, std=0.05) + sinkhorn_iters, norm_eps, hc_eps, sinkhorn_eps, post_mult, layer_input_eps = ( + 3, + 2.5e-3, + 0.25, + 0.05, + 1.5, + 1e-6, + ) + right = (sinkhorn_iters, norm_eps, hc_eps, sinkhorn_eps, post_mult, layer_input_eps) + with validating(op): + out = op(*args, *right, 256) + _compare(out, op.reference(*args, *right, 256), fused_norm=True) + swapped = { + "norm_eps and layer_input_eps": ( + sinkhorn_iters, + layer_input_eps, + hc_eps, + sinkhorn_eps, + post_mult, + norm_eps, + ), + "hc_eps and sinkhorn_eps": ( + sinkhorn_iters, + norm_eps, + sinkhorn_eps, + hc_eps, + post_mult, + layer_input_eps, + ), + "post_mult": (sinkhorn_iters, norm_eps, hc_eps, sinkhorn_eps, 2.0, layer_input_eps), + "sinkhorn_iters": (20, norm_eps, hc_eps, sinkhorn_eps, post_mult, layer_input_eps), + } + for name, wrong in swapped.items(): + assert not _matches(out, op.reference(*args, *wrong, 256), fused_norm=True), ( + f"the band cannot tell {name} apart" + ) + + +def test_cuda_graph_capture_and_replay() -> None: + """The decode step's first pre-mapping is captured: replay on new streams and products.""" + args = _inputs(4, norm=True, seed=31) + scalars(op.CELLS[0].spec) + block_size = big_fuse_block_size(4) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(*args, block_size) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(*args, block_size) + for captured, fresh in zip(args[:6], _inputs(4, norm=True, seed=32)): + captured.copy_(fresh) + graph.replay() + torch.cuda.synchronize() + _compare(out, op.reference(*args, block_size), fused_norm=True) + + +def test_guard_refuses_what_the_kernel_misreads() -> None: + """The kernel reads the product row-major and is compiled for 4 streams.""" + mixes, sqrsum, residual, scale, base, weight = _inputs(8, norm=True, seed=7) + rest = [scale, base, weight, *scalars(op.CELLS[0].spec), 256] + cases = { + "mixes transposed in memory": [mixes.t().contiguous().t(), sqrsum, residual, *rest], + "mixes stored [24, M]": [mixes.t().contiguous(), sqrsum, residual, *rest], + "residual as 2 streams": [mixes, sqrsum, residual.view(8, 2, -1), *rest], + } + for name, args in cases.items(): + with pytest.raises(AssertionError): + with validating(op): + op(*args) + pytest.fail(f"is_valid accepted {name}") diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/mhc/test_modeling_v2_mhc_fused_hc.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/mhc/test_modeling_v2_mhc_fused_hc.py new file mode 100644 index 000000000000..ccf4fde17528 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/mhc/test_modeling_v2_mhc_fused_hc.py @@ -0,0 +1,253 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm5_next fp8 tp4 mhc_fused_hc entry. + +Each cell runs every tactic ``fused_hc_tactic`` can return at its row count on this device -- the +autotuner's candidates, the cache-miss fallback and ``DETERMINISTIC_TACTIC`` -- against one +reference. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.mhc.mhc_fused_hc import ( # noqa: E501 + DETERMINISTIC_TACTIC, + HIDDEN, + default_fused_hc_tactic, + fused_hc_mma_enabled, + fused_hc_scratch, + fused_hc_tactic, + fused_hc_tactics, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.mhc.mhc_fused_hc import ( # noqa: E501 + mhc_fused_hc as op, +) +from tensorrt_llm._torch.autotuner import AutoTuner, autotune +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = [".", "../../../.."] +from _glm53_mhc_inputs import ( # noqa: E402 + block_output, + coefficients, + generator, + mixer, + norm_weight, + scalars, + streams, +) +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 100, reason="the glm5_next fp8 tp4 entries are certified on sm_100" +) + +#: Absolute band of the FP32 coefficients; see `_compare`. +COEFF_ATOL = 5e-3 +_CUDA = torch.device("cuda") + + +def _compare(out, ref) -> None: + """The streams within 3 ulp per element of the row scale and 0.5 ulp relative RMS; ``post`` + and ``comb`` within 5e-3 absolute; the layer input within 4 ulp per element and 1 ulp RMS. + + The streams are the same FP32 sums on both sides, rounded to BF16 once, so they differ only + in rare one-step rounding flips (at most 2 ulp of the row scale). The coefficients differ by + more: the FMA backends feed the mixer the FP32 streams before their BF16 rounding, while the + reference (like the tensor-core backends and the unfused path) uses the stored BF16 streams. + Each stream element then differs by up to half a BF16 step, which moves a mixer output by + about 0.15% of its spread; a CPU transcription of that difference at these inputs measured + at most 6e-4 on ``pre``, 8e-4 on ``post`` and 1.5e-3 on ``comb`` over 3,328 rows, against a + 5e-3 band. The layer input is rounded to BF16 at the same places on both sides, but a ``pre`` + off by 6e-4 flips some of its first rounding: the same transcription measured 2.6 ulp per + element and 0.35 ulp RMS, with the bound at two BF16 steps (4 ulp). A transposed ``comb``, a + lost hidden split or a wrong mixer moves every output by tens to hundreds of times its band. + """ + residual_out, post, comb, layer_input = out + ref_residual, ref_post, ref_comb, ref_layer_input = ref + assert post.shape == ref_post.shape and comb.shape == ref_comb.shape + assert_within_ulp(residual_out, ref_residual, element_ulp=3.0, rms_ulp=0.5) + torch.testing.assert_close(post, ref_post, rtol=0.0, atol=COEFF_ATOL) + torch.testing.assert_close(comb, ref_comb, rtol=0.0, atol=COEFF_ATOL) + assert_within_ulp(layer_input, ref_layer_input, element_ulp=4.0, rms_ulp=1.0) + + +def _matches(out, ref) -> bool: + try: + _compare(out, ref) + except AssertionError: + return False + return True + + +def _inputs(rows: int, norm: bool, seed: int, std: float = 1.0) -> list: + """``[x, residual, post, comb, fn, scale, base, norm_weight]`` of one boundary.""" + gen = generator(seed) + post, comb = coefficients(rows, gen) + fn, scale, base = mixer(gen) + x, residual = block_output(rows, gen, std), streams(rows, gen, std) + return [x, residual, post, comb, fn, scale, base, norm_weight(gen) if norm else None] + + +def _tactics(rows: int) -> list: + """Every tactic ``fused_hc_tactic`` can return at ``rows`` rows on this device.""" + mma = fused_hc_mma_enabled() + sm103 = torch.cuda.get_device_capability() == (10, 3) + tactics = dict.fromkeys(fused_hc_tactics(rows, HIDDEN, mma, sm103)) + tactics[default_fused_hc_tactic(rows, HIDDEN, mma)] = None + tactics[DETERMINISTIC_TACTIC] = None + return list(tactics) + + +def _run(args: list, tactic): + return op(*args, tactic, *fused_hc_scratch(args[0].shape[0], tactic, _CUDA)) + + +def _reference(args: list, tactic=DETERMINISTIC_TACTIC): + return op.reference(*args, tactic, *fused_hc_scratch(args[0].shape[0], tactic, _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) + spec = cell.spec + rows = spec["rows"] + args = _inputs(rows, spec["norm"], seed=zlib.crc32(cell.why.encode()) % 2**31) + args += scalars(spec) + held = [t.clone() for t in args[:4]] + ref = _reference(args) + for tactic in _tactics(rows): + out = _run(args, tactic) + assert out[0].shape == args[1].shape and out[3].shape == args[0].shape + try: + _compare(out, ref) + except AssertionError as e: + raise AssertionError(f"tactic {tactic}: {e}") from e + for name, before, after in zip(("x", "residual", "post", "comb"), held, args[:4]): + assert torch.equal(before, after), f"the op wrote its input {name}" + + +def test_scalars_reach_their_kernel_parameters() -> None: + """Distinct, exaggerated scalars on small streams, where every epsilon moves the outputs. + + GLM-5.3-Flash's two epsilon pairs are equal and too small to move unit-scale streams, so the + cells cannot tell a swapped pair from a right one. Here the streams' mean square is about + 4e-3, next to ``norm_eps``; the references with a pair swapped must fail the band, and the op + must pass it, on the fallback tactic of each backend family. + """ + rows = 64 + args = _inputs(rows, norm=True, seed=11, std=0.05) + right = [3, 2.5e-3, 0.25, 0.05, 1.5, 1e-6] + wrong = { + "norm_eps and layer_input_eps": [3, 1e-6, 0.25, 0.05, 1.5, 2.5e-3], + "hc_eps and sinkhorn_eps": [3, 2.5e-3, 0.05, 0.25, 1.5, 1e-6], + "post_mult": [3, 2.5e-3, 0.25, 0.05, 2.0, 1e-6], + "sinkhorn_iters": [20, 2.5e-3, 0.25, 0.05, 1.5, 1e-6], + } + tactics = dict.fromkeys( + (default_fused_hc_tactic(rows, HIDDEN, fused_hc_mma_enabled()), DETERMINISTIC_TACTIC) + ) + for tactic in tactics: + with validating(op): + out = _run(args + right, tactic) + _compare(out, _reference(args + right)) + for name, scalars_ in wrong.items(): + assert not _matches(out, _reference(args + scalars_)), ( + f"the band cannot tell {name} apart" + ) + + +@pytest.mark.parametrize("all_in_one", [False, True], ids=["fallback", "all_in_one_fma"]) +def test_cuda_graph_capture_and_replay(all_in_one: bool) -> None: + """A captured decode boundary replays on new contents of its inputs and workspaces. + + The all-in-one FMA kernel zeroes its accumulators and counters inside the captured launch. + """ + spec = op.CELLS[1].spec + rows = spec["rows"] + if all_in_one: + tactic = ("fused_all_fma", 2, 1, 0, 1) + else: + tactic = default_fused_hc_tactic(rows, HIDDEN, fused_hc_mma_enabled()) + args = _inputs(rows, spec["norm"], seed=31) + scalars(spec) + scratch = fused_hc_scratch(rows, tactic, _CUDA) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(*args, tactic, *scratch) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(*args, tactic, *scratch) + for seed in (32, 33): + for captured, fresh in zip(args[:8], _inputs(rows, spec["norm"], seed=seed)): + captured.copy_(fresh) + graph.replay() + torch.cuda.synchronize() + _compare(out, _reference(args)) + + +def test_tuned_tactic_is_a_candidate_and_holds(monkeypatch) -> None: + """A cold tuner gives the fallback; a tuning pass picks a candidate that is then looked up.""" + monkeypatch.delenv("FORCE_DETERMINISTIC", raising=False) + monkeypatch.delenv("FORCE_MHC_KERNEL_DETERMINISTIC", raising=False) + spec = op.CELLS[1].spec + rows = spec["rows"] + args = _inputs(rows, spec["norm"], seed=5) + scalars(spec) + tuner = AutoTuner.get() + tuner.clear_cache() + try: + assert fused_hc_tactic(*args) == default_fused_hc_tactic( + rows, HIDDEN, fused_hc_mma_enabled() + ) + with autotune(): + tuned = fused_hc_tactic(*args) + assert tuned in _tactics(rows), f"the tuner returned {tuned}, not a candidate" + assert fused_hc_tactic(*args) == tuned, "the tuned tactic was not looked up" + with validating(op): + out = _run(args, tuned) + _compare(out, _reference(args)) + finally: + tuner.clear_cache() + + +@pytest.mark.parametrize("variable", ["FORCE_DETERMINISTIC", "FORCE_MHC_KERNEL_DETERMINISTIC"]) +def test_force_deterministic_pins_the_tactic(monkeypatch, variable: str) -> None: + """Either switch the built-in mHC path honours pins the tactic, without consulting the tuner.""" + monkeypatch.delenv("FORCE_DETERMINISTIC", raising=False) + monkeypatch.delenv("FORCE_MHC_KERNEL_DETERMINISTIC", raising=False) + monkeypatch.setenv(variable, "1") + spec = op.CELLS[3].spec + args = _inputs(spec["rows"], spec["norm"], seed=3) + scalars(spec) + assert fused_hc_tactic(*args) == DETERMINISTIC_TACTIC + + +def test_guard_refuses_what_the_kernel_misreads() -> None: + """Workspaces the op does not check, and operands it reads as dense rows.""" + spec = op.CELLS[2].spec + rows = spec["rows"] + args = _inputs(rows, spec["norm"], seed=7) + scalars(spec) + all_fma = ("fused_all_fma", 2, 1, 0, 1) + all_mma = ("fused_all_mma", 0, 16, 0, 1) + y, r, done = fused_hc_scratch(rows, all_fma, _CUDA) + _, _, mma_done = fused_hc_scratch(rows, all_mma, _CUDA) + my, mr, _ = fused_hc_scratch(rows, all_mma, _CUDA) + + def with_(index: int, value) -> list: + changed = list(args) + changed[index] = value + return changed + + cases = { + "a short counter array, all-in-one FMA": [*args, all_fma, y, r, done[:1]], + "short accumulators, all-in-one FMA": [*args, all_fma, y[:1], r[:1], done], + "no counter, all-in-one TF32": [*args, all_mma, my, mr, mma_done[:0]], + "comb transposed": [*with_(3, args[3].transpose(1, 2)), all_fma, y, r, done], + "fn of another width": [*with_(4, args[4][:, :HIDDEN].contiguous()), all_fma, y, r, done], + } + for name, call in cases.items(): + with pytest.raises(AssertionError): + with validating(op): + op(*call) + pytest.fail(f"is_valid accepted {name}") diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/mhc/test_modeling_v2_mhc_gemm_sqrsum.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/mhc/test_modeling_v2_mhc_gemm_sqrsum.py new file mode 100644 index 000000000000..59b9aca0f48b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/mhc/test_modeling_v2_mhc_gemm_sqrsum.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm5_next fp8 tp4 mhc_gemm_sqrsum entry.""" + +import zlib + +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.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.mhc.mhc_gemm_sqrsum import ( # noqa: E501 + TILE_N_OPTIONS, + gemm_sqrsum_tile_n, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.mhc.mhc_gemm_sqrsum import ( # noqa: E501 + mhc_gemm_sqrsum as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = [".", "../../../.."] +from _glm53_mhc_inputs import generator, mixer, streams # noqa: E402 +from _validating import certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 100, reason="the glm5_next fp8 tp4 entries are certified on sm_100" +) + +#: The band, as a fraction of the summed magnitudes; see `_compare`. +BAND = 2.0**-18 + + +def _compare(out, ref, x: torch.Tensor, fn: torch.Tensor) -> None: + """Every output within ``2^-18`` of the sum of its terms' magnitudes. + + Both sides add the same terms -- a BF16 by FP32 product is exact in FP32 -- in different + orders: the kernel in a 64-term FMA chain per thread, then a warp butterfly and a sum over 8 + warps; the reference in the FP32 GEMM's and the reduction's own orders. The rounding error of a + sum is a few FP32 ulp (2^-24) of the sum of its terms' magnitudes in a typical draw, and grows + with the chain length only in the worst case; 2^-18 is 64 such ulp, which covers both orders' + tails over every output of a cell. For the sum of squares that magnitude is the result itself. + A lost 1024-element step moves a mixer output by about 32 of its typical terms, some 2^-8 of + the magnitude sum of 16384 terms: about a thousand times the band. + """ + mixes, sqrsum = out + ref_mixes, ref_sqrsum = ref + with true_fp32_matmul(): + magnitude = x.float().abs() @ fn.float().abs().t() + excess = ((mixes - ref_mixes).abs() / (BAND * magnitude)).max().item() + assert excess <= 1.0, f"mixes off by {excess:.2f}x the band" + excess = ((sqrsum - ref_sqrsum).abs() / (BAND * ref_sqrsum)).max().item() + assert excess <= 1.0, f"sqrsum off by {excess:.2f}x the band" + + +def _inputs(rows: int, seed: int): + gen = generator(seed) + residual = streams(rows, gen) + fn, _, _ = mixer(gen) + return residual.view(rows, -1), fn + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + """Every tile width, against the reference; the tile width moves no bit.""" + certifying(op) + rows = cell.spec["rows"] + x, fn = _inputs(rows, seed=zlib.crc32(cell.why.encode()) % 2**31) + ref = op.reference(x, fn, gemm_sqrsum_tile_n(rows)) + first = None + for tile_n in TILE_N_OPTIONS: + out = op(x, fn, tile_n) + assert out[0].shape == (rows, 24) and out[1].shape == (rows,) + _compare(out, ref, x, fn) + if first is None: + first = out + else: + assert torch.equal(out[0], first[0]) and torch.equal(out[1], first[1]), ( + f"tile_n={tile_n} changed the sums" + ) + + +def test_cuda_graph_capture_and_replay() -> None: + """The decode step's first pre-mapping is captured: replay on new streams.""" + x, fn = _inputs(4, seed=31) + tile_n = gemm_sqrsum_tile_n(4) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(x, fn, tile_n) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(x, fn, tile_n) + x.copy_(_inputs(4, seed=32)[0]) + graph.replay() + torch.cuda.synchronize() + _compare(out, op.reference(x, fn, tile_n), x, fn) + + +def test_guard_refuses_what_the_kernel_misreads() -> None: + """The kernel reads both operands with a row stride of K, whatever their strides.""" + x, fn = _inputs(4, seed=7) + wide = torch.zeros(4, 2 * x.shape[1], dtype=x.dtype, device="cuda") + cases = { + "fn transposed in memory": (x, fn.t().contiguous().t()), + "x a column slice": (wide[:, : x.shape[1]], fn), + "fn of another width": (x, fn[:, : x.shape[1] // 2].contiguous()), + } + for name, args in cases.items(): + with pytest.raises(AssertionError): + with validating(op): + op(*args, 1) + pytest.fail(f"is_valid accepted {name}") diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/mhc/test_modeling_v2_mhc_post_mapping.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/mhc/test_modeling_v2_mhc_post_mapping.py new file mode 100644 index 000000000000..3a6fc7fccc92 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/mhc/test_modeling_v2_mhc_post_mapping.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm5_next fp8 tp4 mhc_post_mapping entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.mhc.mhc_post_mapping import ( # noqa: E501 + mhc_post_mapping as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = [".", "../../../.."] +from _glm53_mhc_inputs import block_output, coefficients, generator, streams # noqa: E402 +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 100, reason="the glm5_next fp8 tp4 entries are certified on sm_100" +) + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """3 ulp per element against the row scale, 0.5 ulp relative RMS. + + Kernel and reference form the same FP32 products and sums in the same order; the kernel fuses + each multiply-add, so the two FP32 values differ by about one FP32 ulp, and their single BF16 + rounding differs only where that difference crosses a rounding boundary: rarely, and then by + one BF16 step of the element, at most 2 ulp of its row's scale. A misread stream or a + transposed ``comb`` is off by the streams' own size, hundreds of ulp. + """ + assert_within_ulp(out, ref, element_ulp=3.0, rms_ulp=0.5) + + +def _inputs(rows: int, seed: int): + gen = generator(seed) + post, comb = coefficients(rows, gen) + return block_output(rows, gen), streams(rows, gen), post, comb + + +@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, residual, post, comb = _inputs(cell.spec["rows"], seed=zlib.crc32(cell.why.encode()) % 2**31) + out = op(x, residual, post, comb) + assert out.shape == residual.shape and out.dtype == torch.bfloat16 + assert out.data_ptr() != residual.data_ptr(), "the op must return new streams" + _compare(out, op.reference(x, residual, post, comb)) + # The model keeps post as [M, 4, 1]: the same bytes. + assert torch.equal(op(x, residual, post.unsqueeze(-1), comb), out) + + +def test_cuda_graph_capture_and_replay() -> None: + """The last boundary of a captured decode step replays on new contents of its inputs.""" + x, residual, post, comb = _inputs(4, seed=31) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(x, residual, post, comb) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(x, residual, post, comb) + for captured, fresh in zip((x, residual, post, comb), _inputs(4, seed=32)): + captured.copy_(fresh) + graph.replay() + torch.cuda.synchronize() + _compare(out, op.reference(x, residual, post, comb)) + + +def test_guard_refuses_what_the_kernel_misreads() -> None: + """Strided coefficients are read as dense rows, and a launch of no rows fails later.""" + x, residual, post, comb = _inputs(8, seed=7) + wide = torch.zeros(8, 2 * 4, device="cuda") + cases = { + "post a column slice": (x, residual, wide[:, :4], comb), + "comb transposed": (x, residual, post, comb.transpose(1, 2)), + "no rows": (x[:0], residual[:0], post[:0], comb[:0]), + } + for name, args in cases.items(): + with pytest.raises(AssertionError): + with validating(op): + op(*args) + pytest.fail(f"is_valid accepted {name}") diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/moe/test_modeling_v2_decode_moe.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/moe/test_modeling_v2_decode_moe.py new file mode 100644 index 000000000000..320e0b57dbbf --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/moe/test_modeling_v2_decode_moe.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm5_next fp8 tp4 decode_moe entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4 import ( # noqa: E501 + native_ops, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.moe.decode_moe import ( # noqa: E501 + decode_moe as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 100, reason="the glm5_next fp8 tp4 entries are certified on sm_100" +) + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Both sides dequantize the same FP8 weights exactly and keep the activations in FP32 up to + the single BF16 rounding of the output; they differ in the order of the FP32 sums over 4096 + and 512 products (FMA tiles against torch's matmul), ~1e-6 of the row scale, plus the + reference's FP32 SiLU against the kernel's exp-based one, ~1e-7 relative. 4 ulp element-wise + and 2 ulp RMS admit rare one-ulp flips; a wrong expert, slot weight or clamp is off by the + size of an expert's contribution, many ulp.""" + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0) + + +def _build(spec: dict, seed: int): + """One rank's experts at moe_tp4 shapes, with scales and weights giving O(1) outputs. + + Logit gaps far above FP32 noise keep the top-8 unambiguous, so the reference's + ``torch.topk`` and the kernel's tie rule pick the same experts. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + m, e, h, i = spec["m"], spec["e"], spec["h"], spec["i"] + + def fp8(*shape): + return (torch.randn(*shape, generator=g, device="cuda") * 0.25).to(torch.float8_e4m3fn) + + def scales(*shape, fan_in): + return torch.rand(*shape, generator=g, device="cuda") * fan_in**-0.5 + 0.5 * fan_in**-0.5 + + x = torch.randn(m, h, generator=g, device="cuda").to(torch.bfloat16) + logits = torch.randperm(e, generator=g, device="cuda").float().repeat(m, 1) * 0.01 + logits += torch.arange(m, device="cuda")[:, None] * 0.37 + bias = torch.zeros(e, device="cuda") + w31, s31 = fp8(e, 2 * i, h), scales(e, 2 * i // 128, h // 128, fan_in=h) + w2, s2 = fp8(e, h, i), scales(e, h // 128, i // 128, fan_in=i) + sw31, ss31 = fp8(2 * i, h), scales(2 * i // 128, h // 128, fan_in=h) + sw2, ss2 = fp8(h, i), scales(h // 128, i // 128, fan_in=i) + return (x, logits, bias, 2.5, 10.0, w31, s31, w2, s2, sw31, ss31, sw2, ss2) + + +@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) + _compare(op(*args), op.reference(*args)) + + +def test_clamp_is_applied() -> None: + """A tiny SwiGLU limit makes the clamped and unclamped products disagree; the op must match + the clamped reference.""" + args = list(_build(op.CELLS[1].spec, seed=7)) + args[4] = 0.05 + _compare(op(*args), op.reference(*args)) + + +def test_independent_rows() -> None: + """Rows routed on their own logits with a nonzero correction bias: the step reads ~30 distinct + experts (the kernels stream each once), and the bias picks the experts but not their weights.""" + spec = op.CELLS[1].spec + args = list(_build(spec, seed=13)) + g = torch.Generator(device="cuda").manual_seed(14) + args[1] = torch.randn(spec["m"], spec["e"], generator=g, device="cuda") + args[2] = (torch.rand(spec["e"], generator=g, device="cuda") - 0.5) * 0.1 + _compare(op(*args), op.reference(*args)) + + +def test_schema_and_fake(monkeypatch) -> None: + """The kernel writes only ``out``, as the schema declares, and the fake agrees with the op.""" + args = _build(op.CELLS[0].spec, seed=3) + out = torch.empty(args[0].shape, dtype=torch.bfloat16, device="cuda") + op_args = (*args, out, False) + # SchemaCheckMode compares every input with allclose, which has no float8 CUDA kernel: compare + # the float8 weights as FP32 (the op itself still receives float8). + allclose = torch.allclose + + def allclose_with_float8(actual, expected, *comparison_args, **comparison_kwargs): + if actual.dtype == expected.dtype == torch.float8_e4m3fn: + actual, expected = actual.float(), expected.float() + return allclose(actual, expected, *comparison_args, **comparison_kwargs) + + with monkeypatch.context() as patch: + patch.setattr(torch, "allclose", allclose_with_float8) + torch.library.opcheck(native_ops.decode_moe.default, op_args, test_utils=("test_schema",)) + torch.library.opcheck(native_ops.decode_moe.default, op_args, test_utils=("test_faketensor",)) + + +def test_rejects_operands_the_kernel_cannot_read() -> None: + """The binding refuses, before the launch, operands the kernel would read with the wrong shape, + dtype or alignment.""" + args = _build(op.CELLS[0].spec, seed=5) + rows, hidden = args[0].shape + out = torch.empty(rows, hidden, dtype=torch.bfloat16, device="cuda") + shifted = torch.empty(rows * hidden + 8, dtype=torch.bfloat16, device="cuda") + bad = ( + (0, "x", args[0][:, : hidden // 2].contiguous()), + (0, "x", shifted[1 : 1 + rows * hidden].view(rows, hidden)), + (1, "logits", args[1].to(torch.bfloat16)), + (5, "w31", args[5][:, :512]), + (8, "s2", args[8].transpose(1, 2).contiguous()), + (12, "ss2", args[12].t()), + ) + for index, name, tensor in bad: + call = list(args) + call[index] = tensor + with pytest.raises(RuntimeError, match=f"decode_moe: {name} must"): + native_ops.decode_moe(*call, out, False) + + +def test_cuda_graph_capture_and_replay() -> None: + spec = op.CELLS[1].spec + args = _build(spec, seed=31) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(*args) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(*args) + fresh = _build(spec, seed=32) + args[0].copy_(fresh[0]) + args[1].copy_(fresh[1]) + graph.replay() + torch.cuda.synchronize() + _compare(out, op.reference(*args)) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/moe/test_modeling_v2_glm53_fp8_fp8_block_scale_moe_runner.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/moe/test_modeling_v2_glm53_fp8_fp8_block_scale_moe_runner.py new file mode 100644 index 000000000000..66656c88150c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/moe/test_modeling_v2_glm53_fp8_fp8_block_scale_moe_runner.py @@ -0,0 +1,331 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the fp8_block_scale_moe_runner operation of the glm5_next FP8 target. + +The entry owns the contract: `CELLS`, `reference`, `is_valid`. This file turns a cell into +operands -- GLM-5.3-Flash's 288-expert stack drawn directly in e4m3 / fp32, so both sides read +bit-identical values, and router logits whose top-8 cannot tie -- and holds what no cell +carries: that the cell data can see the correction bias, the clamp and the unbiased combine +weights; the transposed activation-scale hazard at T == H/128; and the guards. +""" + +import zlib +from collections.abc import Callable, Iterator + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.moe.fp8_block_scale_moe_runner import ( # noqa: E501 + BLOCK, + ROUTING_DEEPSEEK_V3, + ROUTING_RENORMALIZE, + Fp8BlockScaleMoeRunner, + route_deepseek_v3, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.moe.fp8_block_scale_moe_runner import ( # noqa: E501 + fp8_block_scale_moe_runner as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +pytestmark = pytest.mark.skipif( + get_sm_version() != 100, reason="fp8_block_scale_moe_runner is certified on sm_100" +) + +DEV = "cuda" +ULP = 2.0**-8 +GLM = Fp8BlockScaleMoeRunner.GLM53_FLASH_TP4 + + +def _ulp(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + """(max, RMS) deviation in bf16 ulp of each reference row's largest magnitude.""" + o, r = out.float(), ref.float() + d = (o - r) / (r.abs().amax(dim=-1, keepdim=True).clamp(min=1e-6) * ULP) + return d.abs().max().item(), d.pow(2).mean().sqrt().item() + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """16 ulp element-wise, 0.75 ulp RMS, both in bf16 ulp of each row's max. + + Each element is normalized by its reference row's largest magnitude, `d = (out - ref) / + (rowmax * 2^-8)`, and gated on `max|d|` and on `sqrt(mean(d^2))`. + + Kernel and reference read bit-identical operands, pick the same experts (the routing below + has no near-ties) and model every rounding the kernel makes: the e4m3 requantization of + the FC1 output and of the SwiGLU product, the bf16 FC2 rows and the bf16 combine weights. + What is left is the order of the fp32 sums, which now and then lands an FC1 output or a + SwiGLU product on the other side of an e4m3 rounding midpoint -- one e4m3 step in one + intermediate element, carried through FC2 into its row -- and the exponentials of the + SiLU and the routing sigmoid, a few fp32 ulp, which can move a bf16 combine weight by one + ulp. Both grow the maximum with the row count and leave the RMS well under one ulp. The + same gate holds the pre-routed path of this kernel at another geometry (H 2560, I 640, + 512 experts, sm_103) with margin, worst seen 10.0 / 0.34; it has not been measured at + GLM-5.3-Flash's geometry on sm_100. + + Wrong references are far outside it: `test_cell_data_see_the_bias_the_clamp_and_the_weights` + asserts that dropping the correction bias from the selection, dropping the clamp, or + weighting by the biased scores each fails this gate on a certified cell. + """ + assert out.dtype == ref.dtype == torch.bfloat16, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + elt, rms = _ulp(out, ref) + assert elt <= 16.0, f"{elt:.2f} ulp element-wise > 16" + assert rms <= 0.75, f"{rms:.3f} ulp RMS > 0.75" + + +# ── operands ────────────────────────────────────────────────────────────── + + +def _draw_stack(hidden: int, inter: int, experts: int, seed: int) -> tuple[torch.Tensor, ...]: + """`(w1, w1s, w2, w2s)` drawn directly in e4m3 / fp32, so every value is exact. + + Codes ~ N(0, 1) in e4m3. FC1 scales in [0.04, 0.08) put FC1 outputs at about 8 RMS from + the activations below, so GLM's limit of 10 clamps a tenth of the gates and a fifth of the + ups; FC2 scales in [0.008, 0.016) bring the layer back to O(10). Drawn 16 experts at a time + to cap the fp32 transient. + """ + g = torch.Generator(device=DEV).manual_seed(seed) + w1 = torch.empty(experts, 2 * inter, hidden, dtype=torch.float8_e4m3fn, device=DEV) + w2 = torch.empty(experts, hidden, inter, dtype=torch.float8_e4m3fn, device=DEV) + for lo in range(0, experts, 16): + hi = min(lo + 16, experts) + w1[lo:hi] = torch.randn(hi - lo, 2 * inter, hidden, device=DEV, generator=g) + w2[lo:hi] = torch.randn(hi - lo, hidden, inter, device=DEV, generator=g) + w1s = 0.04 + 0.04 * torch.rand( + experts, 2 * inter // BLOCK, hidden // BLOCK, device=DEV, generator=g + ) + w2s = 0.008 + 0.008 * torch.rand( + experts, hidden // BLOCK, inter // BLOCK, device=DEV, generator=g + ) + return w1, w1s, w2, w2s + + +@pytest.fixture(scope="module") +def stack() -> Iterator[Callable[..., tuple[torch.Tensor, ...]]]: + """The expert stack for a geometry, drawn once per module and released at teardown. + GLM's is 1.8 GB of e4m3 codes and every GLM cell reads it.""" + built: dict = {} + + def get(hidden: int, inter: int, experts: int) -> tuple[torch.Tensor, ...]: + key = (hidden, inter, experts) + if key not in built: + built[key] = _draw_stack(hidden, inter, experts, seed=hidden + inter + experts) + return built[key] + + try: + yield get + finally: + built.clear() + + +def _routing(tokens: int, experts: int, g: torch.Generator) -> tuple[torch.Tensor, torch.Tensor]: + """FP32 router logits and correction bias whose top-8 cannot tie. + + Each row's biased scores `sigmoid(logit) + bias` are a random permutation of an evenly + spaced ladder from 0.85 down to 0.15: neighbouring ranks sit 0.7 / (experts - 1), 2.4e-3 + at 288 experts, apart -- four orders of magnitude above the fp32 error of the sigmoid and + the bias add -- so the kernel and the reference pick the same experts. The bias, uniform + in +-0.1, moves each unbiased score `ladder - bias` by up to 40 ladder steps, so a + selection that ignored it, or weights that kept it, would differ. + """ + bias = (torch.rand(experts, generator=g, device=DEV) - 0.5) * 0.2 + ladder = torch.linspace(0.85, 0.15, experts, device=DEV) + order = torch.rand(tokens, experts, generator=g, device=DEV).argsort(dim=-1) + choice = torch.empty(tokens, experts, device=DEV).scatter_(1, order, ladder.repeat(tokens, 1)) + score = choice - bias + logits = torch.log(score) - torch.log1p(-score) + return logits.contiguous(), bias.contiguous() + + +def _args(spec: dict, stack: Callable, seed: int) -> dict: + """The op's arguments by name, as the TRTLLM MoE backend passes them for GLM-5.3-Flash: + router logits and bias in, routing in the kernel.""" + hidden, inter, experts = spec["hidden"], spec["inter"], spec["num_experts"] + w1, w1s, w2, w2s = stack(hidden, inter, experts) + g = torch.Generator(device=DEV).manual_seed(seed) + tokens = spec["tokens"] + hs = (2.0 * torch.randn(tokens, hidden, device=DEV, generator=g)).to(torch.float8_e4m3fn) + # [H/128, T]: the K-block-major layout fp8_quantize_1x128 returns here. + hs_scale = 0.5 + torch.rand(hidden // BLOCK, tokens, device=DEV, generator=g) + logits, bias = _routing(tokens, experts, g) + return dict( + routing_logits=logits, + routing_bias=bias, + hidden_states=hs, + hidden_states_scale=hs_scale, + gemm1_weights=w1, + gemm1_weights_scale=w1s, + gemm2_weights=w2, + gemm2_weights_scale=w2s, + num_experts=experts, + top_k=spec["top_k"], + num_fused_shared_experts=0, + n_group=spec["n_group"], + topk_group=spec["topk_group"], + intermediate_size=inter, + local_expert_offset=spec["local_expert_offset"], + local_num_experts=spec["local_num_experts"], + routed_scaling_factor=spec["routed_scaling_factor"], + routing_method_type=ROUTING_DEEPSEEK_V3, + gemm1_clamp_limit=spec["gemm1_clamp_limit"], + ) + + +def _seed(cell: Cell) -> int: + return zlib.crc32(cell.why.encode()) % 2**31 + + +# ── the certified cells ─────────────────────────────────────────────────── + +# The shapes the target calls this operation with, on top of the operation's own `CELLS`. +# +# GLM-5.3-Flash at TP4 hands this op the generation steps its 4-row decode path does not take: +# MTP-3 verifies of 4 rows per request, 2 to 64 requests, and any step of 5 rows or more. +CELLS: tuple[Cell, ...] = op.CELLS + ( + Cell( + why="5 rows: the smallest step past the 4-row decode path; small-token routing", + spec=dict(tokens=5, **GLM), + ), + Cell( + why="MTP-3 verify of 2 requests: 8 rows, the small-token routing kernel's last count", + spec=dict(tokens=8, **GLM), + ), + Cell( + why="MTP-3 verify of 8 requests: 32 rows == H/128, the activation scales are square", + spec=dict(tokens=32, **GLM), + ), + Cell( + why="MTP-3 verify of 64 requests: 256 rows, the single-cluster routing kernel's last", + spec=dict(tokens=256, **GLM), + ), +) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) +def test_certified_cells(cell, certifying, stack): # noqa: F811 — pytest resolves the fixture + certifying(op) + args = _args(cell.spec, stack, seed=_seed(cell)) + out = op(**args) + assert out.shape == args["hidden_states"].shape and out.dtype == torch.bfloat16 + ref = op.reference(**args) + elt, rms = _ulp(out, ref) + print(f" {cell.why[:44]}: {elt:.2f} ulp element-wise, {rms:.3f} ulp RMS") + _compare(out, ref) + + +# ── what a cell cannot carry ────────────────────────────────────────────── + + +def test_cell_data_see_the_bias_the_clamp_and_the_weights(stack): + """The cells certify the bias-steered selection, the clamp and the unbiased weights only + if their data can tell them apart: each wrong reading fails the gate on a certified cell. + """ + args = _args(dict(tokens=8, **GLM), stack, seed=101) + out = op(**args) + _compare(out, op.reference(**args)) + + ids, _ = route_deepseek_v3( + args["routing_logits"], args["routing_bias"], GLM["top_k"], GLM["routed_scaling_factor"] + ) + choice = 0.5 * torch.tanh(0.5 * args["routing_logits"]) + 0.5 + args["routing_bias"] + biased = choice.gather(1, ids) + biased = biased * GLM["routed_scaling_factor"] / biased.sum(dim=-1, keepdim=True) + pre_routed = dict(args, routing_logits=None, routing_bias=None) + wrong = { + "selection without the bias": dict( + args, routing_bias=torch.zeros_like(args["routing_bias"]) + ), + "no clamp": dict(args, gemm1_clamp_limit=None), + "weights from the biased scores": dict( + pre_routed, topk_ids=ids.to(torch.int32), topk_weights=biased.to(torch.bfloat16) + ), + } + for name, wrong_args in wrong.items(): + elt, rms = _ulp(out, op.reference(**wrong_args)) + print(f" {name}: {elt:.2f} ulp element-wise, {rms:.3f} ulp RMS") + with pytest.raises(AssertionError): + _compare(out, op.reference(**wrong_args)) + + +def test_a_transposed_activation_scale_is_silently_wrong(stack): + """At T == H/128 (32 rows at H 4096) the token-major scale buffer has the K-block-major + shape: the op takes it, `is_valid` cannot see it, and the answer is wrong. This is why the + entry says to pass `fp8_quantize_1x128`'s scales through unchanged.""" + args = _args(dict(tokens=32, **GLM), stack, seed=103) + scale = args["hidden_states_scale"] + assert scale.shape[0] == scale.shape[1] + transposed = dict(args, hidden_states_scale=scale.t().contiguous()) + with validating(op): + out = op(**transposed) + with pytest.raises(AssertionError): + _compare(out, op.reference(**args)) + + +# ── the guards ──────────────────────────────────────────────────────────── + + +def _small_args() -> dict: + """A tiny layer for the guards, which refuse before the kernel runs.""" + g = torch.Generator(device=DEV).manual_seed(7) + tokens, hidden, inter, experts = 4, 256, 128, 8 + return dict( + routing_logits=torch.randn(tokens, experts, device=DEV, generator=g), + routing_bias=torch.zeros(experts, device=DEV), + hidden_states=torch.randn(tokens, hidden, device=DEV, generator=g).to(torch.float8_e4m3fn), + hidden_states_scale=torch.ones(hidden // BLOCK, tokens, device=DEV), + gemm1_weights=torch.zeros(experts, 2 * inter, hidden, device=DEV).to(torch.float8_e4m3fn), + gemm1_weights_scale=torch.ones(experts, 2 * inter // BLOCK, hidden // BLOCK, device=DEV), + gemm2_weights=torch.zeros(experts, hidden, inter, device=DEV).to(torch.float8_e4m3fn), + gemm2_weights_scale=torch.ones(experts, hidden // BLOCK, inter // BLOCK, device=DEV), + num_experts=experts, + top_k=2, + num_fused_shared_experts=0, + n_group=1, + topk_group=1, + intermediate_size=inter, + local_expert_offset=0, + local_num_experts=experts, + routed_scaling_factor=1.0, + routing_method_type=ROUTING_DEEPSEEK_V3, + gemm1_clamp_limit=10.0, + ) + + +def test_guard_refuses_deepseek_routing_without_a_bias(): + """The routing kernels read the bias of every expert with no null check: refused before + the kernel could fault the context.""" + with pytest.raises(AssertionError, match="faults on None"): + with validating(op): + op(**dict(_small_args(), routing_bias=None)) + + +def test_guard_refuses_a_bias_the_routing_method_ignores(): + with pytest.raises(AssertionError, match="ignores routing_bias"): + with validating(op): + op(**dict(_small_args(), routing_method_type=ROUTING_RENORMALIZE)) + + +def test_guard_refuses_logits_alongside_topk_ids(): + a = _small_args() + ids = torch.tensor([[0, 1]] * 4, dtype=torch.int32, device=DEV) + wts = torch.full((4, 2), 0.5, dtype=torch.bfloat16, device=DEV) + with pytest.raises(AssertionError, match="ignored when topk_ids"): + with validating(op): + op(**dict(a, routing_bias=None, topk_ids=ids, topk_weights=wts)) + + +def test_guard_refuses_a_strided_hidden_states(): + a = _small_args() + wide = torch.zeros(4, 512, device=DEV).to(torch.float8_e4m3fn) + with pytest.raises(AssertionError, match="hidden_states must be contiguous"): + with validating(op): + op(**dict(a, hidden_states=wide[:, ::2])) + + +def test_guard_refuses_a_window_outside_the_routing_space(): + with pytest.raises(AssertionError, match="outside the routing space"): + with validating(op): + op(**dict(_small_args(), local_expert_offset=8, local_num_experts=8)) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/moe/test_modeling_v2_prefill_moe.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/moe/test_modeling_v2_prefill_moe.py new file mode 100644 index 000000000000..9c8820e96a19 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/moe/test_modeling_v2_prefill_moe.py @@ -0,0 +1,187 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm5_next fp8 tp4 prefill_moe entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.moe.prefill_moe import ( # noqa: E501 + PrefillMoeWorkspace, + route, + route_reference, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.moe.prefill_moe import ( # noqa: E501 + prefill_moe 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() != 100, reason="the glm5_next fp8 tp4 entries are certified on sm_100" +) + +#: GLM-5.3-Flash at moe_tp4. +_GLM = dict(e=288, h=4096, i=512) + +#: The routing weights: the kernel's sigmoid (``0.5 * tanh(0.5 * x) + 0.5`` with libdevice's +#: tanhf) and divisions are within a few FP32 ulp of torch's, and a weight is a ratio of their +#: sums: well under 1e-5 relative. A wrong expert, normalization or scaling factor is off by O(1). +_WEIGHT_RTOL = 1e-5 + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Both sides take the same FP8 inputs and weights, round the FC1 output to BF16 and + requantize the intermediate 1 x 128 to FP8. They differ in the order of the FP32 sums of + the GEMMs (about 1e-7 relative), so a BF16 FC1 output rounds the other way at a rate of + about 1e-4 and then moves an FP8 code of the intermediate one step (12.5% of it) about + once in 30: a relative L2 error of a few 1e-4 on the routed output. A wrong expert, slot + weight, clamp or scale is off by an expert's whole contribution, O(10%) of a row.""" + err = (out.float() - ref.float()).norm() / ref.float().norm() + assert err < 1e-3, f"relative L2 error {err:.2e}" + assert torch.isfinite(out).all() + + +def _build(spec: dict, seed: int) -> list: + """One rank's experts at moe_tp4 shapes and FP8 rows with 1 x 128 scales, O(1) outputs. + + The routing choices ``sigmoid(logit) + bias`` are a per-row permutation of a grid 1.4e-3 + apart, so the top-8 and their order are the same however the sigmoid rounds: the + reference's sort and the kernel's selection pick the same experts in the same slots. + The addend is the shared expert's output in the model. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + t, e, h, i = spec["t"], spec["e"], spec["h"], spec["i"] + + def fp8(*shape: int) -> torch.Tensor: + return (torch.randn(*shape, generator=g, device="cuda") * 0.25).to(torch.float8_e4m3fn) + + def scales(*shape: int, fan_in: int) -> torch.Tensor: + return torch.rand(*shape, generator=g, device="cuda") * fan_in**-0.5 + 0.5 * fan_in**-0.5 + + x = torch.randn(t, h, generator=g, device="cuda").to(torch.bfloat16) + x_fp8, x_scale = torch.ops.trtllm.fp8_quantize_1x128(x) + bias = (torch.rand(e, generator=g, device="cuda") - 0.5) * 0.2 + grid = torch.linspace(0.3, 0.7, e, device="cuda") + order = torch.rand(t, e, generator=g, device="cuda").argsort(dim=1) + logits = torch.logit(grid[order] - bias).contiguous() + w31, s31 = fp8(e, 2 * i, h), scales(e, 2 * i // 128, h // 128, fan_in=h) + w2, s2 = fp8(e, h, i), scales(e, h // 128, i // 128, fan_in=i) + addend = (torch.randn(t, h, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + return [ + x_fp8, + x_scale, + logits, + bias, + 2.5, + 10.0, + w31, + s31, + w2, + s2, + addend, + PrefillMoeWorkspace(), + ] + + +def _check_routing(args: list) -> None: + """The selected experts in slot order exactly, and their weights within ``_WEIGHT_RTOL``.""" + logits, bias, scaling = args[2], args[3], args[4] + ids, wts = route(logits, bias, scaling) + ref_ids, ref_wts = route_reference(logits, bias, scaling) + assert torch.equal(ids, ref_ids), "a row routes to other experts or slots" + torch.testing.assert_close(wts, ref_wts, rtol=_WEIGHT_RTOL, atol=0.0, equal_nan=True) + + +def _check(args: list) -> None: + """The routing, then the routed experts alone (the addend zeroed: the shared expert's + output is about 13 times the routed one and would hide a lost pair), then the sum with the + addend.""" + _check_routing(args) + routed = list(args) + routed[10] = torch.zeros_like(args[10]) + _compare(op(*routed), op.reference(*routed)) + _compare(op(*args), op.reference(*args)) + + +@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) + _check(_build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31)) + + +def test_clamp_is_applied() -> None: + """A tiny SwiGLU limit makes the clamped and unclamped products disagree; the op must match + the clamped reference.""" + args = _build(dict(t=2048, **_GLM), seed=7) + args[5] = 0.05 + _check(args) + + +def test_skewed_routing() -> None: + """Most rows on a few experts: segments of many tiles next to empty experts.""" + args = _build(dict(t=2048, **_GLM), seed=11) + args[2][:, :4] += 100.0 + _check(args) + + +def test_workspace_serves_a_smaller_step() -> None: + """A step after a larger one runs on dense views of the larger buffers.""" + workspace = PrefillMoeWorkspace() + for t, seed in ((2048, 21), (37, 22)): + args = _build(dict(t=t, **_GLM), seed=seed) + args[11] = workspace + with validating(op): + out = op(*args) + _compare(out, op.reference(*args)) + + +def test_one_row_scale_from_the_gather() -> None: + """A one-token prefill: the target all-gathers the scales row-major (a fresh ``[1, K / 128]`` + tensor) and transposes them back, which leaves the size-1 column's stride at K / 128.""" + args = _build(dict(t=1, **_GLM), seed=31) + gathered = args[1].t().clone(memory_format=torch.contiguous_format) + args[1] = gathered.t().contiguous() + assert args[1].shape[1] == 1 and args[1].stride(1) != 1 + with validating(op): + out = op(*args) + _compare(out, op.reference(*args)) + + +def test_nan_logits_route_in_range() -> None: + """A row of NaN router logits must not index past the experts (a device-side assert on every + rank): it routes to experts 0 to 7 with NaN weights, so its output is NaN. A row with some + NaN logits routes among its numbers. Every other row is unaffected.""" + args = _build(dict(t=37, **_GLM), seed=23) + logits = args[2] + logits[5] = float("nan") + logits[6, ::2] = float("nan") + ids, wts = route(logits, args[3], args[4]) + assert torch.equal(ids[5], torch.arange(8, dtype=torch.int32, device="cuda")) + assert bool(torch.isnan(wts[5]).all()) + assert bool((ids[6] % 2 == 1).all()) and bool(torch.isfinite(wts[6]).all()) + _check_routing(args) + out, ref = op(*args), op.reference(*args) + torch.cuda.synchronize() + assert bool(torch.isnan(out[5]).all()) and bool(torch.isnan(ref[5]).all()) + keep = torch.ones(out.shape[0], dtype=torch.bool, device="cuda") + keep[5] = False + _compare(out[keep], ref[keep]) + # The NaN row left nothing behind that breaks the next call. + clean = _build(dict(t=37, **_GLM), seed=24) + _compare(op(*clean), op.reference(*clean)) + + +def test_guard_refuses_mismatched_expert_counts() -> None: + """Expert ids come from the logits' width: a bias, weight or scale of fewer experts is + indexed past its end.""" + args = _build(dict(t=37, **_GLM), seed=25) + for index in (3, 6, 7, 8, 9): + bad = list(args) + bad[index] = args[index][:-1] + with pytest.raises(AssertionError): + with validating(op): + op(*bad) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/norm/test_modeling_v2_glm53_fp8_rms_norm_gated_token_major.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/norm/test_modeling_v2_glm53_fp8_rms_norm_gated_token_major.py new file mode 100644 index 000000000000..7a5711ac0dd0 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/norm/test_modeling_v2_glm53_fp8_rms_norm_gated_token_major.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the norm/rms_norm_gated_token_major operation of the glm5_next FP8 target. + +The gate is a column section of a wider projection output, as the KDA layer reads it: ``z[t]`` is token t's +``[heads, N]`` block of the output-gate GEMM, and ``x`` holds the recurrence's per-head rows. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.norm.rms_norm_gated_token_major import ( # noqa: E501 + rms_norm_gated_token_major as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 -- path above; fixture by name + +pytestmark = pytest.mark.skipif(get_sm_version() != 100, reason="certified on sm_100") + +# The shapes the shipped target normalizes, by target directory under `models/`: GLM-5.3-Flash's KDA output +# gate at TP4 (16 of 64 heads, head_dim 128, rms_norm_eps 1e-5, sigmoid gate), on the prefill path (decode and +# verify fuse the norm into their kernels). +CELLS: tuple[Cell, ...] = op.CELLS + ( + Cell( + why="a 13,312-token prompt: 212,992 rows, more programs than SMs (no PDL)", + spec=dict(tokens=13312, heads=16, width=128, eps=1e-5, gate_activation="sigmoid"), + ), + Cell( + why="a 32,768-token chunked-prefill step", + spec=dict(tokens=32768, heads=16, width=128, eps=1e-5, gate_activation="sigmoid"), + ), + Cell( + why="a 17-token prompt: fewer programs than SMs, so the PDL launch", + spec=dict(tokens=17, heads=16, width=128, eps=1e-5, gate_activation="sigmoid"), + ), +) + + +def _build(spec: dict, seed: int): + """``(x, z, weight, eps, fp8_scale, gate_activation)``: z is the second half of a ``[2, tokens, heads * N]`` + batched GEMM output (the forget gate is the first), so its token stride is ``heads * N`` and it starts + ``tokens * heads * N`` elements into the buffer.""" + gen = torch.Generator(device="cuda").manual_seed(seed) + tokens, heads, width = spec["tokens"], spec["heads"], spec["width"] + x = torch.randn(tokens * heads, width, generator=gen, device="cuda").to(torch.bfloat16) + gates = torch.randn(2, tokens, heads * width, generator=gen, device="cuda").to(torch.bfloat16) + z = gates[1].view(tokens, heads, width) + weight = (1 + 0.1 * torch.randn(width, generator=gen, device="cuda")).to(torch.bfloat16) + fp8_scale = None + if "fp8_scale" in spec: + fp8_scale = torch.tensor([spec["fp8_scale"]], dtype=torch.float32, device="cuda") + return x, z, weight, spec["eps"], fp8_scale, spec["gate_activation"] + + +def _assert_matches(out: torch.Tensor, ref: torch.Tensor, fp8: bool) -> None: + """Both sides compute in fp32 and round once, so they differ by the kernel's reduction order and its + sigmoid / sqrt implementations (relative 1e-6), which can move the rounded output by one ulp of bf16 + (2^-8 to 2^-7 relative): ``rtol`` is two such ulps (2 * bf16 eps), ``atol`` 1e-5 covers outputs near zero. + With fp8 the same one-ulp bf16 difference before quantization can move an e4m3 code by one step (2^-3 + relative): ``rtol`` 2^-3 plus a bf16 ulp, ``atol`` one e4m3 subnormal step (2^-9) for codes near zero.""" + if fp8: + torch.testing.assert_close(out.float(), ref.float(), rtol=2.0**-3 + 2.0**-8, atol=2.0**-9) + else: + torch.testing.assert_close( + out.float(), ref.float(), rtol=2 * torch.finfo(torch.bfloat16).eps, atol=1e-5 + ) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 -- pytest resolves the fixture by this name + certifying(op) + args = _build(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + x, z = args[0].clone(), args[1].clone() + out = op(*args) + fp8 = args[4] is not None + assert out.shape == x.shape and out.is_contiguous() + assert out.dtype == (torch.float8_e4m3fn if fp8 else x.dtype) + assert torch.equal(args[0], x) and torch.equal(args[1], z), "the op mutated its inputs" + _assert_matches(out, op.reference(*args), fp8) + + +def test_cuda_graph_replay() -> None: + """Compiled outside capture, the launch replays rewritten inputs bit for bit (the PDL form: 8 tokens).""" + x, z, weight, eps, _, gate = _build( + dict(tokens=8, heads=16, width=128, eps=1e-5, gate_activation="sigmoid"), seed=4 + ) + op(x, z, weight, eps, None, gate) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(x, z, weight, eps, None, gate) + for _ in range(3): + x.copy_(torch.randn_like(x)) + z.copy_(torch.randn_like(z)) + graph.replay() + assert torch.equal(out, op(x, z, weight, eps, None, gate)) + + +def test_guard_refuses_a_head_strided_gate() -> None: + """A gate whose heads are not packed per token (here every head padded to 2N) is read at ``h * N`` anyway: + the kernel returns a plausible norm gated by the wrong columns.""" + x, _, weight, eps, _, gate = _build( + dict(tokens=16, heads=6, width=128, eps=1e-5, gate_activation="sigmoid"), seed=5 + ) + z = torch.randn(16, 6, 256, device="cuda").to(torch.bfloat16)[:, :, :128] + with pytest.raises(AssertionError, match="dense per token"): + with validating(op): + op(x, z, weight, eps, None, gate) + + +def test_guard_refuses_an_unknown_gate() -> None: + """Any gate name but 'sigmoid' compiles the silu gate: 'relu' would return a silu-gated norm.""" + args = list( + _build(dict(tokens=4, heads=16, width=128, eps=1e-5, gate_activation="silu"), seed=6) + ) + args[5] = "relu" + with pytest.raises(AssertionError, match="gate_activation must be"): + with validating(op): + op(*args) + + +def test_guard_refuses_per_column_fp8_scale() -> None: + """The kernel reads only the first element of the scale: a per-column scale quantizes every column by it.""" + args = list( + _build(dict(tokens=8, heads=16, width=128, eps=1e-5, gate_activation="sigmoid"), seed=7) + ) + args[4] = torch.linspace(0.5, 3.0, 128, dtype=torch.float32, device="cuda") + with pytest.raises(AssertionError, match="exactly one element"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/quantization/test_modeling_v2_glm53_fp8_fp8_quantize_1x128.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/quantization/test_modeling_v2_glm53_fp8_fp8_quantize_1x128.py new file mode 100644 index 000000000000..de96b4b2f26c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/quantization/test_modeling_v2_glm53_fp8_fp8_quantize_1x128.py @@ -0,0 +1,169 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the fp8_quantize_1x128 operation of the glm5_next FP8 target. + +The entry owns the contract. This file turns a cell into an input, holds every code +disagreement of a cell to a rounding tie, and adds what a cell cannot carry: the loud +rejection of a strided input. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.quantization.fp8_quantize_1x128 import ( # noqa: E501 + BLOCK, + block_scales, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.quantization.fp8_quantize_1x128 import ( # noqa: E501 + fp8_quantize_1x128 as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +pytestmark = pytest.mark.skipif( + get_sm_version() != 100, reason="fp8_quantize_1x128 is certified on sm_100" +) + +DEV = "cuda" + +#: The default build divides exactly, so its scales are the reference's. A `FAST_MATH=ON` build +#: divides approximately, each division within CUDA's documented 2 ulp, and the returned scale is +#: two divisions deep: the gate admits that build. +SCALE_ULP_GATE = 4 + +#: A code can only disagree with the exactly rounded one when the scale it multiplies by is off, +#: so when the exact product sits within that division's 2 ulp of an e4m3 rounding midpoint. +TIE_REL_GATE = 2.0 * 2.0**-23 + + +def _ordinal(codes: torch.Tensor) -> torch.Tensor: + """Sign-magnitude ordinal of an e4m3 byte: adjacent codes differ by 1.""" + b = codes.view(torch.uint8).int() + return torch.where(b >= 128, -(b & 0x7F), b & 0x7F) + + +def _compare( + out: tuple[torch.Tensor, torch.Tensor], + ref: tuple[torch.Tensor, torch.Tensor], +) -> None: + """Scales within 4 fp32 ulp; codes at most one e4m3 step apart, >= 99.5% bit-exact. + + The default build's divisions are exact, so it matches the reference bit for bit. The + bands admit a `FAST_MATH=ON` build, whose approximate divisions can put a scale up to 2 ulp + from `amax / 448` per division and a code on the other side of a rounding midpoint when + the exact product sits within that band of it -- one e4m3 step, never more. + `_assert_misses_are_ties` holds every such miss to the band. bf16 inputs make exact + midpoints common (`x * 448 / amax` of two bf16 values often has few significant bits), + which is why that build's bit-exact floor is 99.5% and not higher. + + A wrong recipe is far outside this: a power-of-two scale, a per-row or per-tensor amax, + blocks down the row axis or offset by 64 columns, or 240 as the e4m3 maximum each change + most codes and move the scales by O(1) relative. + """ + data, scales = out + ref_data, ref_scales = ref + assert data.dtype == ref_data.dtype == torch.float8_e4m3fn, (data.dtype, ref_data.dtype) + assert scales.dtype == ref_scales.dtype == torch.float32, (scales.dtype, ref_scales.dtype) + assert data.shape == ref_data.shape, (data.shape, ref_data.shape) + assert scales.shape == ref_scales.shape, (scales.shape, ref_scales.shape) + ulp = ( + scales.contiguous().view(torch.int32).long() + - ref_scales.contiguous().view(torch.int32).long() + ) + worst = int(ulp.abs().max()) + assert worst <= SCALE_ULP_GATE, f"scale {worst} fp32 ulp off > {SCALE_ULP_GATE}" + step = (_ordinal(data) - _ordinal(ref_data)).abs() + assert int(step.max()) <= 1, f"a code is {int(step.max())} e4m3 steps off" + exact = 1.0 - (step != 0).float().mean().item() + assert exact >= 0.995, f"only {100 * exact:.4f}% of codes bit-exact" + + +def _input(rows: int, cols: int, seed: int) -> torch.Tensor: + g = torch.Generator(device=DEV).manual_seed(seed) + return (torch.randn(rows, cols, device=DEV, generator=g) * 3.0).to(torch.bfloat16) + + +def _assert_misses_are_ties(x: torch.Tensor, data: torch.Tensor, ref_data: torch.Tensor) -> float: + """Every code disagreement is a rounding tie: the exact product sits within the scale's + 2-ulp band of the midpoint between the two codes.""" + qs, _ = block_scales(x, False) + prod = x.float() * qs.repeat_interleave(BLOCK, dim=1)[:, : x.shape[1]] + diff = data.float() != ref_data.float() # -0 and +0 are the same code here + if not bool(diff.any()): + return 0.0 + mid = ((data.float() + ref_data.float()) / 2).double() + rel = ((prod.double()[diff] - mid[diff]).abs() / mid[diff].abs()).max().item() + assert rel <= TIE_REL_GATE, f"a miss {rel:.3e} from its midpoint is not a tie" + return rel + + +# The shapes the target calls this operation with, on top of the operation's own `CELLS`. +# +# GLM-5.3-Flash at TP4: every FFN input is the hidden row (4096); the down projections read +# the SwiGLU output at the per-rank intermediate (dense 12288 / 4 = 3072, shared expert +# 2048 / 4 = 512); the sequence-parallel MLA prefill projects q_resid (q_lora 1536) through q_b. +# Generation steps are MTP-3 verifies of 4 rows per request (up to 64 requests); a 13,312-row +# prefill runs sequence parallel, each rank quantizing its own quarter (3,328 rows) of the +# FFN input and the down-projection inputs over every gathered row. +CELLS: tuple[Cell, ...] = op.CELLS + ( + Cell( + why="MTP-3 verify of one request: 4 rows at hidden 4096", + spec=dict(rows=4, cols=4096, use_ue8m0=False), + ), + Cell( + why="MTP-3 verify of 64 requests: 256 rows at hidden 4096", + spec=dict(rows=256, cols=4096, use_ue8m0=False), + ), + Cell( + why="sequence-parallel prefill: one rank's 3,328 rows of 13,312 at hidden 4096", + spec=dict(rows=3328, cols=4096, use_ue8m0=False), + ), + Cell( + why="dense MLP down_proj input at TP4 (3072), 256 rows", + spec=dict(rows=256, cols=3072, use_ue8m0=False), + ), + Cell( + why="dense MLP down_proj input over a 13,312-row prefill step", + spec=dict(rows=13312, cols=3072, use_ue8m0=False), + ), + Cell( + why="shared expert down_proj input at TP4 (512), 256 rows", + spec=dict(rows=256, cols=512, use_ue8m0=False), + ), + Cell( + why="shared expert down_proj input over a 13,312-row prefill step", + spec=dict(rows=13312, cols=512, use_ue8m0=False), + ), + Cell( + why="MLA q_b input (q_lora 1536) over one rank's 3,328 sequence-parallel rows", + spec=dict(rows=3328, cols=1536, use_ue8m0=False), + ), +) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + spec = cell.spec + x = _input(spec["rows"], spec["cols"], seed=zlib.crc32(cell.why.encode()) % 2**31) + data, scales = op(x, spec["use_ue8m0"]) + nb = (spec["cols"] + BLOCK - 1) // BLOCK + assert data.shape == x.shape and data.is_contiguous() and data.device == x.device + assert scales.shape == (nb, spec["rows"]) and scales.is_contiguous() + ref = op.reference(x, spec["use_ue8m0"]) + _compare((data, scales), ref) + rel = _assert_misses_are_ties(x, data, ref[0]) + exact = (data.view(torch.uint8) == ref[0].view(torch.uint8)).float().mean().item() + print(f" {cell.why[:44]}: {100 * exact:.4f}% bit-exact, worst tie {rel:.3e}") + + +def test_rejects_a_strided_input(): + """The op's own precondition, loud, which is why there is no guard for it.""" + x = torch.randn(8, 8192, dtype=torch.bfloat16, device=DEV)[:, :4096] + with pytest.raises((RuntimeError, NotImplementedError)): + op(x) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ssm/test_modeling_v2_glm53_fp8_kda_decode.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ssm/test_modeling_v2_glm53_fp8_kda_decode.py new file mode 100644 index 000000000000..2b4f737c29f0 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ssm/test_modeling_v2_glm53_fp8_kda_decode.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 ssm/kda_decode operation of the glm5_next FP8 target. + +Inputs as the KDA layer hands them over: x_q / x_k / x_v column sections of the fused ``[q | k | v]`` +projection's rows, the forget and output gates the two halves of the batched ``[f_b; g_b]`` GEMM's output, beta +the H-wide tail of the ``[f_a | g_a | b]`` projection's rows, the conv sections views of one packed conv pool. +Two layouts: ``indexed``, the pools' rows by slot (an fp32 recurrent pool), each pool a view into wider per-slot +records as the hybrid cache manager lays it out; ``dense``, no slot indices, request b on row b of fp32 staging +copies (how a bf16 pool's rows are decoded, the pool being bf16 by default). The reference is fp32 torch. +""" + +import math +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.ssm.kda_decode import ( # noqa: E501 + kda_decode as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402,F401 -- fixture by name + assert_within_ulp, + certifying, + validating, +) + +pytestmark = pytest.mark.skipif(get_sm_version() != 100, reason="certified on sm_100") + +#: GLM-5.3-Flash's KDA at TP4: 64 heads over 4 ranks, K = V = 128, conv width 4. +HEADS, HEAD_DIM, WIDTH = 16, 128, 4 +DIM = HEADS * HEAD_DIM +#: gate_lower_bound, the q scale K ** -0.5 and rms_norm_eps (the output norm's epsilon). +LOWER_BOUND, SCALE, ONORM_EPS = -5.0, HEAD_DIM**-0.5, 1e-5 +#: The [f_a | g_a | b] projection's row width at TP4: beta is its last HEADS columns. +BFA_WIDTH = 2 * HEAD_DIM + HEADS +#: Elements past each slot's row in its record (the other layers' rows in the real manager). +SLOT_PAD = 4096 +STATE_ROW = HEADS * HEAD_DIM * HEAD_DIM +CONV_ROW = 3 * DIM * (WIDTH - 1) +#: The fp32 state's band, as a fraction of the used rows' largest entry (see the cells test). +STATE_BAND = 1e-4 + +# The decode steps the shipped target runs, by target directory under `models/`: GLM-5.3-Flash's KDA at TP4, +# CUDA-graph batches up to 64 requests. +CELLS: tuple[Cell, ...] = op.CELLS + ( + Cell( + why="1 request on an fp32 pool by slot: the four-CTA cluster kernel", + spec=dict(batch=1, layout="indexed", pool=72), + ), + Cell( + why="8 requests on an fp32 pool by slot: the legacy compact-heads kernel", + spec=dict(batch=8, layout="indexed", pool=72), + ), + Cell( + why="64 requests on an fp32 pool by slot: the two-stage bulk-copy kernel", + spec=dict(batch=64, layout="indexed", pool=72), + ), + Cell( + why="1 request on staged rows of a bf16 pool (no slot indices)", + spec=dict(batch=1, layout="dense"), + ), + Cell( + why="8 requests on staged rows of a bf16 pool (no slot indices)", + spec=dict(batch=8, layout="dense"), + ), + Cell( + why="64 requests on staged rows of a bf16 pool (no slot indices)", + spec=dict(batch=64, layout="dense"), + ), +) + + +def _heads(cols: torch.Tensor) -> torch.Tensor: + """``[B, H * 128]`` columns as the op's ``[1, B, H, 128]`` view (the row stride kept).""" + return cols.unflatten(-1, (HEADS, HEAD_DIM)).unsqueeze(0) + + +def _build(spec: dict, seed: int): + """``(records, args)``: the pools' backing storage and the op's arguments.""" + gen = torch.Generator(device="cuda").manual_seed(seed) + batch = spec["batch"] + + def rnd(*shape, scale=1.0): + return torch.randn(*shape, generator=gen, device="cuda") * scale + + qkv = rnd(batch, 3 * DIM).to(torch.bfloat16) + gates = rnd(2, batch, DIM).to(torch.bfloat16) + bfa = rnd(batch, BFA_WIDTH).to(torch.bfloat16) + taps = rnd(3, DIM, WIDTH, scale=0.3).to(torch.bfloat16) + w_t = [taps[i].t().contiguous() for i in range(3)] + zeros = torch.zeros(DIM, dtype=torch.bfloat16, device="cuda") + a_log = torch.empty(HEADS, device="cuda").uniform_(1, 16, generator=gen).log() + dt_bias = torch.empty(DIM, device="cuda").uniform_( + math.log(1e-3), math.log(1e-1), generator=gen + ) + onorm_weight = 1 + 0.1 * rnd(HEAD_DIM) + if spec["layout"] == "indexed": + pool = spec["pool"] + conv_record = rnd(pool, CONV_ROW + SLOT_PAD, scale=0.5).to(torch.bfloat16) + state_record = rnd(pool, STATE_ROW + SLOT_PAD, scale=0.05) + conv = conv_record[:, :CONV_ROW].view(pool, 3 * DIM, WIDTH - 1) + state = state_record[:, :STATE_ROW].view(pool, HEADS, HEAD_DIM, HEAD_DIM) + slots = torch.randperm(pool, generator=gen, device="cuda")[:batch].to(torch.int32) + records = [conv_record, state_record] + else: + conv = rnd(batch, 3 * DIM, WIDTH - 1, scale=0.5).to(torch.bfloat16) + state = rnd(batch, HEADS, HEAD_DIM, HEAD_DIM, scale=0.05) + slots = None + records = [conv, state] + output = torch.empty(batch, 1, HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + args = ( + _heads(qkv[:, :DIM]), + _heads(qkv[:, DIM : 2 * DIM]), + _heads(qkv[:, 2 * DIM :]), + *w_t, + zeros, + zeros, + zeros, + conv[:, :DIM], + conv[:, DIM : 2 * DIM], + conv[:, 2 * DIM :], + a_log, + _heads(gates[0]), + dt_bias, + bfa[:, 2 * HEAD_DIM :].unsqueeze(0), + _heads(gates[1]), + onorm_weight, + slots, + state, + LOWER_BOUND, + SCALE, + ONORM_EPS, + output, + ) + return records, args + + +def _assert_state_close(got: torch.Tensor, ref: torch.Tensor) -> None: + scale = ref.abs().max().clamp_min(1e-6) + err = ((got - ref).abs().max() / scale).item() + assert err <= STATE_BAND, f"state error {err:.3e} of its largest entry exceeds {STATE_BAND:.0e}" + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 -- pytest resolves the fixture by this name + """The output, the used state rows and conv windows against the fp32 reference; nothing else written. + + Tolerances. Every product and sum is fp32 on both sides (bf16 inputs, exact bf16 x bf16 products); the + kernel differs only by its fast exp / rcp / rsqrt (a few fp32 ulps each, about 1e-6 relative) and its + reduction order. The output is the gated RMS norm, rounded once to bf16 on both sides, so outputs differ + by at most one bf16 ulp plus about 1e-5 relative: per 128-wide head row that is at most 2 units of 2^-8 of + the row's largest magnitude (one ulp at the row maximum is 2^-7 of it); ``element_ulp=4`` doubles that, + and the relative RMS of two independent half-ulp roundings is about 0.6 of 2^-8 (``rms_ulp=1``). The fp32 + state rows carry the approximate gate and key through one step: about 1e-6 of their largest entry, + ``STATE_BAND`` = 1e-4. The conv windows take the raw bf16 inputs: bit equality. + """ + certifying(op) + records, args = _build(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + before = [r.clone() for r in records] + sections, slots, state, output = args[9:12], args[18], args[19], args[-1] + want_out, want_state, *want_sections = op.reference(*args) + assert op(*args) is output + assert_within_ulp( + output.reshape(-1, HEAD_DIM), want_out.reshape(-1, HEAD_DIM), element_ulp=4, rms_ulp=1 + ) + rows = torch.arange(output.shape[0], device="cuda") if slots is None else slots.long() + _assert_state_close(state[rows], want_state[rows]) + for got, want in zip(sections, want_sections): + assert torch.equal(got, want), "the conv windows are raw-input copies" + if slots is not None: + untouched = [s for s in range(state.shape[0]) if s not in slots.tolist()] + for record, old in zip(records, before): + assert torch.equal(record[untouched], old[untouched]), "the op wrote an unused slot" + for record, old, row in zip(records, before, (CONV_ROW, STATE_ROW)): + assert torch.equal(record[:, row:], old[:, row:]), "the op wrote past a slot's row" + + +def test_cuda_graph_replay() -> None: + """Captured once, the step replays rewritten inputs and slots to the bits of eager steps on a copy of the + pools (nothing allocates or syncs inside the op).""" + records, args = _build(dict(batch=8, layout="indexed", pool=16), seed=21) + copies = [r.clone() for r in records] + copy_args = list(args) + pool = copies[0].shape[0] + copy_args[9:12] = [ + copies[0][:, :CONV_ROW].view(pool, 3 * DIM, WIDTH - 1)[:, i * DIM : (i + 1) * DIM] + for i in range(3) + ] + copy_args[19] = copies[1][:, :STATE_ROW].view(pool, HEADS, HEAD_DIM, HEAD_DIM) + copy_args[-1] = torch.empty_like(args[-1]) + op(*args) # warm-up step, mirrored on the copies + op(*copy_args) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(*args) + inputs = ( + args[0], + args[13], + args[15], + args[16], + ) # x_q (a view of the qkv rows), g, beta, onorm_g + slots = args[18] + gen = torch.Generator(device="cuda").manual_seed(22) + for _ in range(3): + for t in inputs: + t.copy_(torch.randn(t.shape, generator=gen, device="cuda").to(t.dtype)) + slots.copy_( + torch.randperm(pool, generator=gen, device="cuda")[: slots.numel()].to(torch.int32) + ) + graph.replay() + copy_args[-1] = torch.empty_like(args[-1]) + want = op(*copy_args) + assert torch.equal(args[-1], want) + assert all(torch.equal(r, c) for r, c in zip(records, copies)) + + +def test_guard_refuses_a_shared_slot() -> None: + """Two requests on one slot both read and write its state and conv window in one launch: a race.""" + _, args = _build(dict(batch=2, layout="indexed", pool=4), seed=23) + args = list(args) + args[18] = torch.tensor([3, 3], dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match="each row once"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_slot_past_the_pool() -> None: + """The kernel indexes the state and conv pools with the slot unchecked.""" + _, args = _build(dict(batch=1, layout="indexed", pool=4), seed=24) + args = list(args) + args[18] = torch.tensor([4], dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match="lie in the pool"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ssm/test_modeling_v2_glm53_fp8_kda_mtp_decode.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ssm/test_modeling_v2_glm53_fp8_kda_mtp_decode.py new file mode 100644 index 000000000000..435117a88b0f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ssm/test_modeling_v2_glm53_fp8_kda_mtp_decode.py @@ -0,0 +1,264 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the ssm/kda_mtp_decode operation of the glm5_next FP8 target. + +Inputs as the KDA layer hands them over in a verify step: x_q / x_k / x_v column sections of the fused +``[q | k | v]`` projection's packed rows (1 + num_spec per request), the forget and output gates the two halves of +the batched ``[f_b; g_b]`` GEMM's output, beta the H-wide tail of the ``[f_a | g_a | b]`` projection's rows. Every +pool (state, replay conv caches, k / g / v / beta caches) is a view into wider per-slot records, as the hybrid +cache manager lays them out, and the accepted-draft counts are the first int32 of a 16-byte per-slot record. The +replay caches hold what a previous verify would have left: unit-norm keys, log-decays in (lower_bound, 0), SiLU +values and sigmoid betas. The reference is fp32 torch, token by token. +""" + +import math +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.ssm.kda_mtp_decode import ( # noqa: E501 + kda_mtp_decode as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402,F401 -- fixture by name + assert_within_ulp, + certifying, + validating, +) + +pytestmark = pytest.mark.skipif(get_sm_version() != 100, reason="certified on sm_100") + +#: GLM-5.3-Flash's KDA at TP4: 64 heads over 4 ranks, K = V = 128, conv width 4 (a 3-tap window). +HEADS, HEAD_DIM, WINDOW = 16, 128, 3 +DIM = HEADS * HEAD_DIM +#: gate_lower_bound, the q scale K ** -0.5 and rms_norm_eps (the output norm's epsilon). +LOWER_BOUND, SCALE, ONORM_EPS = -5.0, HEAD_DIM**-0.5, 1e-5 +#: The [f_a | g_a | b] projection's row width at TP4: beta is its last HEADS columns. +BFA_WIDTH = 2 * HEAD_DIM + HEADS +#: Elements past each slot's row in its record (the other layers' rows in the real manager). +SLOT_PAD = 1024 +STATE_ROW = HEADS * HEAD_DIM * HEAD_DIM +#: The committed fp32 state's band, as a fraction of the used rows' largest entry (see the cells test). +STATE_BAND = 1e-4 + +# The MTP-3 verify steps the shipped target runs, by target directory under `models/`: GLM-5.3-Flash's KDA at +# TP4, 1 + 3 tokens per request, CUDA-graph batches up to 64 requests, previous rounds accepting 0..3 drafts; +# the pool bf16 (the default recurrent-state dtype) or fp32. +CELLS: tuple[Cell, ...] = op.CELLS + ( + Cell( + why="MTP-3 verify of 1 request, 2 drafts pending, bf16 pool: the cluster-pair kernel", + spec=dict(batch=1, num_spec=3, pending=(2,), state_dtype=torch.bfloat16, pool=8), + ), + Cell( + why="MTP-3 verify of 8 requests, 0..3 drafts pending, bf16 pool: one CTA per (head, request)", + spec=dict( + batch=8, + num_spec=3, + pending=(0, 1, 2, 3, 3, 2, 1, 0), + state_dtype=torch.bfloat16, + pool=12, + ), + ), + Cell( + why="MTP-3 verify of 64 requests, 0..3 drafts pending, bf16 pool: the two-CTA-per-SM kernel", + spec=dict( + batch=64, + num_spec=3, + pending=tuple(n % 4 for n in range(64)), + state_dtype=torch.bfloat16, + pool=72, + ), + ), + Cell( + why="MTP-3 verify of 8 requests on an fp32 pool (mamba_ssm_cache_dtype float32)", + spec=dict( + batch=8, + num_spec=3, + pending=(3, 0, 2, 1, 1, 3, 0, 2), + state_dtype=torch.float32, + pool=12, + ), + ), +) + + +def _heads(cols: torch.Tensor) -> torch.Tensor: + """``[T, H * 128]`` columns as the op's ``[1, T, H, 128]`` view (the row stride kept).""" + return cols.unflatten(-1, (HEADS, HEAD_DIM)).unsqueeze(0) + + +def _records(pool: int, width: int, dtype: torch.dtype, values: torch.Tensor) -> torch.Tensor: + """A ``[pool, width + SLOT_PAD]`` record buffer whose first ``width`` columns hold ``values``.""" + record = torch.full((pool, width + SLOT_PAD), 7.0, dtype=dtype, device="cuda") + record[:, :width] = values.reshape(pool, width).to(dtype) + return record + + +def _build(spec: dict, seed: int): + """``(records, slots, args)``: the per-slot record buffers, the batch's slots and the op's arguments.""" + gen = torch.Generator(device="cuda").manual_seed(seed) + batch, num_spec, pool = spec["batch"], spec["num_spec"], spec["pool"] + tokens = batch * (1 + num_spec) + cols = WINDOW + num_spec + + def rnd(*shape, scale=1.0): + return torch.randn(*shape, generator=gen, device="cuda") * scale + + qkv = rnd(tokens, 3 * DIM).to(torch.bfloat16) + gates = rnd(2, tokens, DIM).to(torch.bfloat16) + bfa = rnd(tokens, BFA_WIDTH).to(torch.bfloat16) + taps = ( + rnd(3, DIM, WINDOW + 1, scale=0.3).to(torch.bfloat16).float() + ) # checkpoint bf16 taps, held fp32 + a_log = torch.empty(HEADS, device="cuda").uniform_(1, 16, generator=gen).log() + dt_bias = torch.empty(DIM, device="cuda").uniform_( + math.log(1e-3), math.log(1e-1), generator=gen + ) + onorm_weight = 1 + 0.1 * rnd(HEAD_DIM) + + conv_records = [ + _records(pool, cols * DIM, torch.bfloat16, rnd(pool, cols, DIM, scale=0.5)) + for _ in range(3) + ] + convs = [r[:, : cols * DIM].view(pool, cols, DIM).transpose(1, 2) for r in conv_records] + state_record = _records(pool, STATE_ROW, spec["state_dtype"], rnd(pool, STATE_ROW, scale=0.05)) + state = state_record[:, :STATE_ROW].view(pool, HEADS, HEAD_DIM, HEAD_DIM) + keys = rnd(pool, num_spec, HEADS, HEAD_DIM) + keys = keys / keys.norm(dim=-1, keepdim=True) + cache_values = ( + (keys, torch.bfloat16, DIM), + (LOWER_BOUND * torch.sigmoid(rnd(pool, num_spec, DIM)), torch.float32, DIM), + (torch.nn.functional.silu(rnd(pool, num_spec, DIM)), torch.bfloat16, DIM), + (torch.sigmoid(rnd(pool, num_spec, HEADS)), torch.bfloat16, HEADS), + ) + cache_records = [ + _records(pool, num_spec * width, dtype, values) for values, dtype, width in cache_values + ] + caches = [ + r[:, : num_spec * width].view(pool, num_spec, width) + for r, (_, _, width) in zip(cache_records, cache_values) + ] + count_record = torch.zeros(pool, 4, dtype=torch.int32, device="cuda") + num_accepted = count_record[:, 0] + slots = torch.randperm(pool, generator=gen, device="cuda")[:batch].to(torch.int32) + num_accepted[slots.long()] = torch.tensor(spec["pending"], dtype=torch.int32, device="cuda") + args = ( + _heads(qkv[:, :DIM]), + _heads(qkv[:, DIM : 2 * DIM]), + _heads(qkv[:, 2 * DIM :]), + *(taps[i].contiguous() for i in range(3)), + *convs, + _heads(gates[0]), + bfa[:, 2 * HEAD_DIM :].unsqueeze(0), + a_log, + dt_bias, + state, + *caches, + slots, + num_accepted, + num_spec, + LOWER_BOUND, + SCALE, + _heads(gates[1]), + onorm_weight, + ONORM_EPS, + ) + records = [*conv_records, state_record, *cache_records, count_record] + return records, slots, args + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 -- pytest resolves the fixture by this name + """The output, the committed state, the conv windows and the replay caches against the fp32 reference; + unused slots, record padding and the accepted counts untouched. + + Tolerances. The recurrence runs on tensor cores in 3xTF32 (within about 2e-6 of fp32); gates, SiLU, norms + and sigmoids use fast exp / rcp / rsqrt (about 1e-6 relative). The output is the gated RMS norm of the raw + output, which the kernel first rounds to bf16 (2^-9 relative, carried through the norm), then rounded once + to bf16 on both sides: at most 1.25 bf16 ulps apart, i.e. 2.5 units of 2^-8 of a head row's largest + magnitude; ``element_ulp=4`` leaves 1.5 units for the rest, and the relative RMS of the three roundings is + about 0.7 of 2^-8 (``rms_ulp=1.5``). The committed state: an fp32 pool within 1e-4 of the used rows' + largest entry (50 times the 3xTF32 error); a bf16 pool is rounded to nearest on both sides from values + that close, so entries differ by at most one bf16 ulp (``rtol`` 2^-7) plus that fp32 band. The replay k / v + / beta caches are bf16 roundings of values within about 1e-6 relative, so at most one ulp apart + (``rtol`` 2^-7, ``atol`` 1e-5 for entries near zero); the fp32 log-decay cache is ``lower_bound`` times an + approximate sigmoid, within about 5e-6 absolute (``atol`` 1e-5). The conv caches are raw-input copies: bit + equality. + """ + certifying(op) + records, slots, args = _build(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + before = [r.clone() for r in records] + x_v, state = args[2], args[13] + want = op.reference(*args) + out = op(*args) + assert out.shape == x_v.shape and out.dtype == torch.bfloat16 and out.is_contiguous() + assert_within_ulp( + out.reshape(-1, HEAD_DIM), want[0].reshape(-1, HEAD_DIM), element_ulp=4, rms_ulp=1.5 + ) + rows = slots.long() + got_state, want_state = state[rows].float(), want[1][rows].float() + fp32_band = STATE_BAND * want_state.abs().max().item() + if state.dtype == torch.float32: + torch.testing.assert_close(got_state, want_state, rtol=0.0, atol=fp32_band) + else: + torch.testing.assert_close(got_state, want_state, rtol=2.0**-7, atol=fp32_band) + for got, expected in zip(args[6:9], want[2:5]): + assert torch.equal(got, expected), "the conv caches are raw-input copies" + k_cache, g_cache, v_cache, beta_cache = args[14:18] + for got, expected in ((k_cache, want[5]), (v_cache, want[7]), (beta_cache, want[8])): + torch.testing.assert_close(got.float(), expected.float(), rtol=2.0**-7, atol=1e-5) + torch.testing.assert_close(g_cache, want[6], rtol=0.0, atol=1e-5) + untouched = [s for s in range(state.shape[0]) if s not in slots.tolist()] + for record, old in zip(records, before): + assert torch.equal(record[untouched], old[untouched]), "the op wrote an unused slot" + padding = [r.shape[1] - SLOT_PAD for r in records[:-1]] + for record, old, width in zip(records, before, padding): + assert torch.equal(record[:, width:], old[:, width:]), "the op wrote past a slot's row" + assert torch.equal(records[-1], before[-1]), ( + "the accepted counts are the runtime's, not the op's" + ) + + +def test_guard_refuses_a_count_past_num_spec() -> None: + """A count above num_spec is clamped by the kernel: the request silently replays fewer drafts than the + runtime accepted.""" + _, slots, args = _build( + dict(batch=2, num_spec=3, pending=(0, 1), state_dtype=torch.float32, pool=4), seed=31 + ) + args[19][slots[1].long()] = 4 + with pytest.raises(AssertionError, match="num_accepted_tokens"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_shared_slot() -> None: + """Two requests on one slot both replay, commit and cache into it in one launch: a race.""" + _, _, args = _build( + dict(batch=2, num_spec=3, pending=(0, 1), state_dtype=torch.float32, pool=4), seed=32 + ) + args = list(args) + args[18] = torch.tensor([2, 2], dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match="each slot once"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_misaligned_operand() -> None: + """The kernels take every token, weight and pool operand as 16-byte aligned without checking: a gate + starting one element into its buffer would be read from the wrong address.""" + _, _, args = _build( + dict(batch=2, num_spec=3, pending=(0, 1), state_dtype=torch.float32, pool=4), seed=33 + ) + args = list(args) + g = args[9] + buffer = torch.empty(g.numel() + 1, dtype=g.dtype, device="cuda") + shifted = buffer[1:].view(g.shape) + shifted.copy_(g) + args[9] = shifted + with pytest.raises(AssertionError, match="16-byte aligned"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ssm/test_modeling_v2_glm53_fp8_kda_prefill.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ssm/test_modeling_v2_glm53_fp8_kda_prefill.py new file mode 100644 index 000000000000..65e41447d0b5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ssm/test_modeling_v2_glm53_fp8_kda_prefill.py @@ -0,0 +1,259 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the ssm/kda_prefill operation of the glm5_next FP8 target. + +Inputs as the KDA layer hands them over: q / k / v the prefill conv's contiguous outputs, the forget gate the +first half of the batched ``[f_b; g_b]`` GEMM's output, beta the H-wide tail of the ``[f_a | g_a | b]`` +projection's rows (row stride 2 * 128 + H), the state pool a view into wider per-slot records (as the hybrid +cache manager lays it out), fresh rows zeroed by the caller, and the scheduler's planned ``seq_order`` / +``cu_chunks``. The reference is the fp32 token-by-token recurrence. +""" + +import itertools +import math +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.ssm.kda_prefill import ( # noqa: E501 + kda_prefill as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 -- path above; fixture by name + +pytestmark = pytest.mark.skipif(get_sm_version() != 100, reason="certified on sm_100") + +HEAD_DIM = 128 +#: GLM-5.3-Flash's linear_attn_config gate_lower_bound, and the q scale K ** -0.5. +LOWER_BOUND, SCALE = -5.0, HEAD_DIM**-0.5 +#: Elements past each slot's state in its record (the other layers' states in the real manager). +SLOT_PAD = 4096 +#: Bands against the fp32 recurrence, as fractions of the compared tensor's largest magnitude (see `_bands`). +OUT_BAND, STATE_BAND = 2e-2, 1e-2 + +# The prefill steps the shipped target runs, by target directory under `models/`: GLM-5.3-Flash's KDA at TP4 +# (16 of 64 heads), the engine route it forces, the scheduler's planned metadata, --max_num_tokens 32768 and +# 13,312-token prompts; the pool bf16 (the default recurrent-state dtype) or fp32. +CELLS: tuple[Cell, ...] = op.CELLS + ( + Cell( + why="one 13,312-token prompt into a fresh slot, bf16 pool", + spec=dict( + heads=16, + seq_lens=(13312,), + slots=(5,), + fresh=(True,), + state_dtype=torch.bfloat16, + mode="engine", + planned=True, + pool=8, + ), + ), + Cell( + why="a 32,768-token step: two whole prompts and a 6,144-token first chunk", + spec=dict( + heads=16, + seq_lens=(13312, 13312, 6144), + slots=(3, 0, 6), + fresh=(True, True, True), + state_dtype=torch.bfloat16, + mode="engine", + planned=True, + pool=8, + ), + ), + Cell( + why="the next step: the 7,168-token continuation from its slot's state, two fresh prompts", + spec=dict( + heads=16, + seq_lens=(7168, 13312, 12288), + slots=(6, 1, 4), + fresh=(False, True, True), + state_dtype=torch.bfloat16, + mode="engine", + planned=True, + pool=8, + ), + ), + Cell( + why="a continuation and a fresh prompt on an fp32 pool (mamba_ssm_cache_dtype float32)", + spec=dict( + heads=16, + seq_lens=(4096, 2048), + slots=(2, 7), + fresh=(False, True), + state_dtype=torch.float32, + mode="engine", + planned=True, + pool=8, + ), + ), +) + + +def _schedule(lens) -> tuple[torch.Tensor, torch.Tensor]: + """``(seq_order, cu_chunks)`` as the scheduler's metadata builds them: rows of one int32 buffer.""" + n = len(lens) + order = sorted(range(n), key=lambda i: (-lens[i], i)) + chunks = [0, *itertools.accumulate((length + 15) // 16 for length in lens)] + schedule = torch.zeros(2, n + 1, dtype=torch.int32, device="cuda") + schedule[0, :n] = torch.tensor(order, dtype=torch.int32, device="cuda") + schedule[1] = torch.tensor(chunks, dtype=torch.int32, device="cuda") + return schedule[0, :n], schedule[1] + + +def _build(spec: dict, seed: int): + """``(record, args)``: the state pool's per-slot records and the op's arguments. + + The value distributions are the Kimi K3 entry's (unit q / k / v / gate / beta logits, A_log = log U(1, 16), + dt_bias = U(log 1e-3, log 1e-1), continuing states of spread 0.1), whose bands this test keeps. + """ + gen = torch.Generator(device="cuda").manual_seed(seed) + heads, lens, pool = spec["heads"], spec["seq_lens"], spec["pool"] + tokens = sum(lens) + shape = (1, tokens, heads, HEAD_DIM) + q, k, v = ( + torch.randn(*shape, generator=gen, device="cuda").to(torch.bfloat16) for _ in range(3) + ) + gates = torch.randn(2, tokens, heads * HEAD_DIM, generator=gen, device="cuda").to( + torch.bfloat16 + ) + g = gates[0].view(shape) + bfa = torch.randn(tokens, 2 * HEAD_DIM + heads, generator=gen, device="cuda").to(torch.bfloat16) + beta = bfa[:, 2 * HEAD_DIM :].unsqueeze(0) + A_log = torch.empty(heads, device="cuda").uniform_(1, 16, generator=gen).log() + dt_bias = torch.empty(heads * HEAD_DIM, device="cuda").uniform_( + math.log(1e-3), math.log(1e-1), generator=gen + ) + row = heads * HEAD_DIM * HEAD_DIM + record = (torch.randn(pool, row + SLOT_PAD, generator=gen, device="cuda") * 0.1).to( + spec["state_dtype"] + ) + state_pool = record[:, :row].view(pool, heads, HEAD_DIM, HEAD_DIM) + for slot, fresh in zip(spec["slots"], spec["fresh"]): + if fresh: + state_pool[slot].zero_() # the caller's reset_recurrent_state_rows + state_indices = torch.tensor(spec["slots"], dtype=torch.int32, device="cuda") + cu_seqlens = torch.tensor([0, *itertools.accumulate(lens)], dtype=torch.int64, device="cuda") + seq_order = cu_chunks = None + max_seqlen = 0 + if spec["planned"]: + seq_order, cu_chunks = _schedule(lens) + max_seqlen = max(lens) + args = ( + q, + k, + v, + g, + beta, + state_pool, + state_indices, + cu_seqlens, + A_log, + dt_bias, + SCALE, + LOWER_BOUND, + spec["mode"], + max_seqlen, + seq_order, + cu_chunks, + 0, + ) + return record, args + + +def _assert_scaled(got: torch.Tensor, ref: torch.Tensor, band: float) -> None: + scale = ref.float().abs().max().clamp_min(1e-6) + err = ((got.float() - ref.float()).abs().max() / scale).item() + assert err <= band, f"max error {err:.3e} of the largest magnitude exceeds {band:.3e}" + + +def _bands(state_dtype: torch.dtype) -> tuple[float, float]: + """``(output band, state band)`` as fractions of the compared tensor's largest magnitude. + + The kernel runs 16-token chunks on tensor cores: q, k, the decayed keys, the state and each chunk's update + enter the MMAs rounded to bf16 (2^-9 relative per operand), with fp32 accumulation and the state carried in + fp32 across the call. Each output and update term thus carries about 2^-8 relative error; over 128 keys of + mixed sign it stays a few 2^-9 of the output's scale, and the delta rule's projection ``I - beta k k^T`` + keeps re-writing the state along each new key, so the state's error does not grow with length. The bands + are the ones the Kimi K3 entry holds for this kernel on the same value distributions: 2e-2 of each + sequence's largest output, 1e-2 of a slot's largest fp32 state entry. A bf16 pool also rounds every final + entry once, at most one bf16 ulp (2^-7 of the entry), which the state band adds. + """ + return OUT_BAND, STATE_BAND + (2.0**-7 if state_dtype == torch.bfloat16 else 0.0) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 -- pytest resolves the fixture by this name + """Each sequence's output and its slot's final state within the bands of `_bands`; every other slot and + all record padding bit-identical.""" + certifying(op) + record, args = _build(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + before = record.clone() + v, state_pool, state_indices, cu_seqlens = args[2], args[5], args[6], args[7] + o_ref, pool_ref = op.reference(*args) + o = op(*args) + assert o.shape == v.shape and o.dtype == torch.bfloat16 and o.is_contiguous() + out_band, state_band = _bands(state_pool.dtype) + bounds = cu_seqlens.tolist() + slots = state_indices.tolist() + for seq, slot in enumerate(slots): + lo, hi = bounds[seq], bounds[seq + 1] + if hi > lo: + _assert_scaled(o[0, lo:hi], o_ref[0, lo:hi], out_band) + _assert_scaled(state_pool[slot], pool_ref[slot], state_band) + untouched = [s for s in range(state_pool.shape[0]) if s not in slots] + assert torch.equal(record[untouched], before[untouched]), "the op wrote an unselected slot" + row = state_pool[0].numel() + assert torch.equal(record[:, row:], before[:, row:]), "the op wrote past a slot's state" + + +def _small(**overrides): + spec = dict( + heads=16, + seq_lens=(37, 100), + slots=(0, 1), + fresh=(False, False), + state_dtype=torch.float32, + mode="engine", + planned=True, + pool=4, + ) + spec.update(overrides) + return _build(spec, seed=9)[1] + + +def test_guard_refuses_a_shared_slot() -> None: + """Two sequences on one slot each read and write it in one launch: a race, no error.""" + args = _small(slots=(1, 1)) + with pytest.raises(AssertionError, match="unique slot"): + with validating(op): + op(*args) + + +def test_guard_refuses_offsets_that_miss_tokens() -> None: + """Offsets ending short of the packed tokens leave the tail rows of ``o`` unwritten.""" + args = list(_small()) + args[7] = torch.tensor([0, 37, 136], dtype=torch.int64, device="cuda") + with pytest.raises(AssertionError, match="partition"): + with validating(op): + op(*args) + + +@pytest.mark.parametrize("bad", ("order", "chunks")) +def test_guard_refuses_stale_planned_metadata(bad) -> None: + """The planned schedule replaces the kernel's own: a stale one (lengths changed, schedule not refreshed) + is taken as is.""" + args = list(_small()) + if bad == "order": + args[14] = torch.tensor([0, 1], dtype=torch.int32, device="cuda") + match = "descending" + else: + args[15] = torch.tensor([0, 3, 9], dtype=torch.int32, device="cuda") + match = "prefix-sum" + with pytest.raises(AssertionError, match=match): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ssm/test_modeling_v2_glm53_fp8_kda_prefill_conv.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ssm/test_modeling_v2_glm53_fp8_kda_prefill_conv.py new file mode 100644 index 000000000000..371288d60373 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ssm/test_modeling_v2_glm53_fp8_kda_prefill_conv.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the ssm/kda_prefill_conv operation of the glm5_next FP8 target. + +The conv pool is a view into wider per-slot records, as the hybrid cache manager lays it out, so a write past a +slot's window shows up in the record's padding; the source is the token-major projection GEMM's output. +""" + +import itertools +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.ssm.kda_prefill_conv import ( # noqa: E501 + kda_prefill_conv as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 -- path above; fixture by name + +pytestmark = pytest.mark.skipif(get_sm_version() != 100, reason="certified on sm_100") + +HEAD_DIM, WIDTH = 128, 4 +#: Elements past each slot's [3D, 3] window in its record (the other layers' windows in the real manager). +SLOT_PAD = 1024 + +# The prefill steps the shipped target runs, by target directory under `models/`: GLM-5.3-Flash's KDA at TP4 +# (16 of 64 heads, [q | k | v] = 6144 channels), with --max_num_tokens 32768 and 13,312-token prompts. +CELLS: tuple[Cell, ...] = op.CELLS + ( + Cell( + why="one 13,312-token prompt into a fresh slot", + spec=dict(heads=16, seq_lens=(13312,), slots=(5,), has_init=(False,), pool=8), + ), + Cell( + why="a 32,768-token step: two whole 13,312-token prompts and a 6,144-token first chunk", + spec=dict( + heads=16, + seq_lens=(13312, 13312, 6144), + slots=(3, 0, 6), + has_init=(False, False, False), + pool=8, + ), + ), + Cell( + why="the next step: the 7,168-token continuation from its slot's window, two fresh prompts", + spec=dict( + heads=16, + seq_lens=(7168, 13312, 12288), + slots=(6, 1, 4), + has_init=(True, False, False), + pool=8, + ), + ), +) + + +def _build(spec: dict, seed: int): + """``(record, args)``: the conv pool's per-slot records and the op's arguments. + + The taps (std 0.5 over four unit inputs) give SiLU arguments of about unit spread: both signs of the + kernel's fast-SiLU range in every cell, and in the long cells also its cancellation range (x < -4, which + the kernel recomputes exactly). + """ + gen = torch.Generator(device="cuda").manual_seed(seed) + heads, lens, pool = spec["heads"], spec["seq_lens"], spec["pool"] + channels = 3 * heads * HEAD_DIM + tokens = sum(lens) + width = spec.get("src_width", channels) + src = torch.randn(tokens, width, generator=gen, device="cuda").to(torch.bfloat16) + weight = (torch.randn(channels, WIDTH, generator=gen, device="cuda") * 0.5).to(torch.bfloat16) + window = channels * (WIDTH - 1) + record = torch.randn(pool, window + SLOT_PAD, generator=gen, device="cuda").to(torch.bfloat16) + conv_state = record[:, :window].view(pool, channels, WIDTH - 1) + starts = torch.tensor([0, *itertools.accumulate(lens)], dtype=torch.int32, device="cuda") + slots = torch.tensor(spec["slots"], dtype=torch.int32, device="cuda") + has_init = torch.tensor(spec["has_init"], dtype=torch.bool, device="cuda") + return record, (src, weight, conv_state, starts, slots, has_init, max(lens)) + + +def _assert_conv_close(out: torch.Tensor, ref: torch.Tensor) -> None: + """The taps accumulate in fp32 on both sides (the kernel with exactly halved weights), so the outputs differ + by the SiLU and the two bf16 roundings. The kernel's SiLU is ``h + h * tanh.approx(h)`` with ``h = x / 2`` + for x >= -4: tanh.approx has relative error at most 2^-10.987, which is at most 2^-11 of the output for + x >= 0 and at most ``2 * tanh(2) * 2^-10.987 < 9.5e-4`` absolute for -4 <= x < 0 (smaller x is recomputed + exactly). The two roundings to bf16 put the results at most one ulp (2^-7 relative) apart. + ``rtol`` = 2^-7 + 2^-10 (the ulp plus twice the relative SiLU bound), ``atol`` = 2^-9 (twice the absolute + SiLU bound).""" + torch.testing.assert_close(out.float(), ref.float(), rtol=2.0**-7 + 2.0**-10, atol=2.0**-9) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 -- pytest resolves the fixture by this name + """q / k / v within the SiLU band; the conv pool bit-identical to the reference (its windows are raw-input + copies, the other slots untouched); no write past a slot's window.""" + certifying(op) + record, args = _build(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + before = record.clone() + src, conv_state = args[0], args[2] + src_before = src.clone() + q_ref, k_ref, v_ref, state_ref = op.reference(*args) + q, k, v = op(*args) + for out, ref in ((q, q_ref), (k, k_ref), (v, v_ref)): + assert out.shape == ref.shape and out.dtype == torch.bfloat16 and out.is_contiguous() + _assert_conv_close(out, ref) + assert torch.equal(conv_state, state_ref) + window = conv_state[0].numel() + assert torch.equal(record[:, window:], before[:, window:]), "the op wrote past a slot's window" + assert torch.equal(src, src_before), "the op mutated its source" + + +def test_guard_refuses_a_shared_slot() -> None: + """Two sequences on one slot each read and write its window in one launch: a race, no error.""" + _, args = _build( + dict(heads=16, seq_lens=(40, 9), slots=(2, 2), has_init=(True, True), pool=4), seed=8 + ) + with pytest.raises(AssertionError, match="unique"): + with validating(op): + op(*args) + + +def test_guard_refuses_offsets_short_of_the_tokens() -> None: + """Rows past ``starts[-1]`` belong to no sequence: the op leaves them uninitialized and returns them.""" + _, args = _build( + dict(heads=16, seq_lens=(40, 9), slots=(0, 1), has_init=(False, True), pool=4), seed=9 + ) + short = torch.tensor([0, 40, 48], dtype=torch.int32, device="cuda") + args = (*args[:3], short, *args[4:]) + with pytest.raises(AssertionError, match="partition"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_slot_outside_the_pool() -> None: + """Slot -1 is the kernel's pass-through: it copies the raw inputs out unconvolved.""" + _, args = _build(dict(heads=16, seq_lens=(16,), slots=(0,), has_init=(False,), pool=4), seed=10) + args = (*args[:4], torch.tensor([-1], dtype=torch.int32, device="cuda"), *args[5:]) + with pytest.raises(AssertionError, match="lie in the pool"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ssm/test_modeling_v2_glm53_fp8_reset_recurrent_state_rows.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ssm/test_modeling_v2_glm53_fp8_reset_recurrent_state_rows.py new file mode 100644 index 000000000000..43b9fe8a224c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/ssm/test_modeling_v2_glm53_fp8_reset_recurrent_state_rows.py @@ -0,0 +1,123 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the ssm/reset_recurrent_state_rows operation of the glm5_next FP8 target. + +The pools are views into wider per-slot records, as the hybrid cache manager lays them out, so a write past a +slot's row shows up in the record's padding. The op writes zeros or nothing, so every comparison is exact. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell +from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.ops.ssm.reset_recurrent_state_rows import ( # noqa: E501 + reset_recurrent_state_rows as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 -- path above; fixture by name + +pytestmark = pytest.mark.skipif(get_sm_version() != 100, reason="certified on sm_100") + +#: GLM-5.3-Flash's KDA at TP4: 64 heads over 4 ranks, K = V = 128. +HEADS, HEAD_DIM = 16, 128 +STATE_ROW = HEADS * HEAD_DIM * HEAD_DIM +#: Elements past each slot's row in its record (the other layers' rows in the real manager). +SLOT_PAD = 4096 + +# The resets the shipped target issues before each KDA prefill, by target directory under `models/`. +CELLS: tuple[Cell, ...] = op.CELLS + ( + Cell( + why="a 13,312-token prompt's prefill: one fresh request on the default bf16 pool", + spec=dict( + slots=8, + indices=(5,), + has_init=(False,), + state_dtype=torch.bfloat16, + ), + ), + Cell( + why="a chunked-prefill step: the continuing chunk keeps its row, two fresh prompts are zeroed", + spec=dict( + slots=8, + indices=(6, 1, 4), + has_init=(True, False, False), + state_dtype=torch.bfloat16, + ), + ), + Cell( + why="the same step on an fp32 pool (kv_cache_config.mamba_ssm_cache_dtype float32)", + spec=dict( + slots=8, + indices=(6, 1, 4), + has_init=(True, False, False), + state_dtype=torch.float32, + ), + ), +) + + +def _build(spec: dict, seed: int) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]: + """``(records, args)``: the per-slot records, nonzero everywhere, and the op's arguments viewing them.""" + gen = torch.Generator(device="cuda").manual_seed(seed) + slots = spec["slots"] + records = torch.randn(slots, STATE_ROW + SLOT_PAD, generator=gen, device="cuda").to( + spec["state_dtype"] + ) + states = records[:, :STATE_ROW].view(slots, HEADS, HEAD_DIM, HEAD_DIM) + indices = torch.tensor( + spec["indices"], dtype=spec.get("index_dtype", torch.int32), device="cuda" + ) + has_init = torch.tensor(spec["has_init"], dtype=torch.bool, device="cuda") + return records, (states, indices, has_init) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 -- pytest resolves the fixture by this name + """The fresh rows are zero, every other row and all record padding bit-identical (the op writes 0.0 or + nothing, so there is no band).""" + certifying(op) + records, args = _build(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + before = records.clone() + want_states = op.reference(*args) + assert op(*args) is None + assert torch.equal(args[0], want_states) + assert torch.equal(records[:, STATE_ROW:], before[:, STATE_ROW:]), ( + "the op wrote past a slot's row" + ) + + +def test_empty_batch_launches_nothing() -> None: + """No request, no launch: the pools are untouched (a zero-row grid is not launched at all).""" + records, args = _build( + dict(slots=4, indices=(), has_init=(), state_dtype=torch.float32), seed=3 + ) + before = records.clone() + op(*args) + assert torch.equal(records, before) + + +def test_guard_refuses_a_flag_per_slot() -> None: + """``has_initial_states`` is read by request index: a per-slot flag vector of another length is taken and + zeroes rows by the wrong requests' flags.""" + _, (states, indices, _) = _build( + dict(slots=8, indices=(6, 1), has_init=(True, False), state_dtype=torch.float32), seed=4 + ) + per_slot = torch.zeros(8, dtype=torch.bool, device="cuda") + with pytest.raises(AssertionError, match="one flag per request"): + with validating(op): + op(states, indices, per_slot) + + +def test_guard_refuses_overlapping_rows() -> None: + """A pool view whose slot stride is shorter than its row makes zeroing one slot clear part of the next.""" + flat = torch.randn(4 * 1024, device="cuda") + overlapping = flat.as_strided((4, 2, 1024), (512, 1024, 1)) + indices = torch.tensor([1], dtype=torch.int32, device="cuda") + has_init = torch.tensor([False], device="cuda") + with pytest.raises(AssertionError, match="must not overlap"): + with validating(op): + op(overlapping, indices, has_init) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/test_modeling_v2_glm53_fp8_fake_ops.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/test_modeling_v2_glm53_fp8_fake_ops.py new file mode 100644 index 000000000000..0d6fe583aee4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/glm_5_3_flash_fp8__sm_10x__tp4/test_modeling_v2_glm53_fp8_fake_ops.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The target's fakes (``fake_ops.py``) against its native ops (``csrc/glm53Fp8Ops.cpp``). + +Every op the target's library defines has a fake, ``fake_``, that takes the op's schema +arguments by name and in order, so an argument added to a schema without its fake fails here +rather than at the first trace. Both ops then run under ``FakeTensorMode``, directly and through +their op wrappers, and the wrappers' outputs have the shapes and dtypes the real ops write. + +Nothing runs on a device: the tensors are fake. The target's library must be installed, though; +without it the module is skipped. +""" + +import importlib +import inspect + +import pytest +import torch +from torch._subclasses.fake_tensor import FakeTensorMode + +from tensorrt_llm._torch._experimental.modeling_v2._native import namespace_of +from tensorrt_llm._torch._experimental.modeling_v2._router_index import NoTarget + +_TARGET = ( + "tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4" +) + +try: + _target = importlib.import_module(_TARGET) +except NoTarget as error: + pytest.skip(f"this installation does not carry the target: {error}", allow_module_level=True) + +native_ops = _target.native_ops +fake_ops = importlib.import_module(f"{_TARGET}.fake_ops") +decode_moe_module = importlib.import_module(f"{_TARGET}.ops.moe.decode_moe") +boundary_module = importlib.import_module(f"{_TARGET}.ops.comm.mhc_ar_boundary") + +# The shapes the kernels are compiled for (csrc/decodeMoe.h, csrc/mhcArBoundary.h). +HIDDEN, INTER, EXPERTS, BLOCK = 4096, 512, 288, 128 +MULT, MIX = boundary_module.MULT, (2 + boundary_module.MULT) * boundary_module.MULT +WORKSPACE_POINTERS = 3 * boundary_module.WORLD + 2 + + +def _target_schemas() -> dict: + """The ops the target's library registers in its namespace, by name.""" + namespace = namespace_of(_TARGET) + return { + schema.name.split("::", 1)[1]: schema + for schema in torch._C._jit_get_all_schemas() + if schema.name.split("::", 1)[0] == namespace + } + + +def test_every_op_has_a_fake_taking_its_schema_arguments(): + schemas = _target_schemas() + assert {"decode_moe", "mhc_ar_boundary"} <= set(schemas), sorted(schemas) + for name, schema in schemas.items(): + fake = getattr(fake_ops, f"fake_{name}", None) + assert fake is not None, f"fake_ops.py has no fake_{name} for {schema}" + expected = [argument.name for argument in schema.arguments] + actual = list(inspect.signature(fake).parameters) + assert actual == expected, ( + f"fake_{name} takes {actual}; the schema in csrc/glm53Fp8Ops.cpp is {schema}" + ) + + +def _empty(*shape, dtype): + return torch.empty(shape, dtype=dtype, device="cuda") + + +def _decode_moe_inputs(rows: int) -> tuple: + """``DecodeMoe.__call__``'s arguments, at the shapes the kernel is compiled for.""" + fp8 = torch.float8_e4m3fn + return ( + _empty(rows, HIDDEN, dtype=torch.bfloat16), + _empty(rows, EXPERTS, dtype=torch.float32), + _empty(EXPERTS, dtype=torch.float32), + 2.5, + 10.0, + _empty(EXPERTS, 2 * INTER, HIDDEN, dtype=fp8), + _empty(EXPERTS, 2 * INTER // BLOCK, HIDDEN // BLOCK, dtype=torch.float32), + _empty(EXPERTS, HIDDEN, INTER, dtype=fp8), + _empty(EXPERTS, HIDDEN // BLOCK, INTER // BLOCK, dtype=torch.float32), + _empty(2 * INTER, HIDDEN, dtype=fp8), + _empty(2 * INTER // BLOCK, HIDDEN // BLOCK, dtype=torch.float32), + _empty(HIDDEN, INTER, dtype=fp8), + _empty(HIDDEN // BLOCK, INTER // BLOCK, dtype=torch.float32), + ) + + +def _boundary_inputs(rows: int) -> tuple: + """``MhcArBoundary.__call__``'s arguments, at the shapes the kernel is compiled for.""" + return ( + _empty(rows, HIDDEN, dtype=torch.bfloat16), + _empty(WORKSPACE_POINTERS, dtype=torch.int64), + 0, + _empty(rows, MULT, HIDDEN, dtype=torch.bfloat16), + _empty(rows, MULT, dtype=torch.float32), + _empty(rows, MULT, MULT, dtype=torch.float32), + _empty(MIX, MULT * HIDDEN, dtype=torch.float32), + _empty(3, dtype=torch.float32), + _empty(MIX, dtype=torch.float32), + _empty(HIDDEN, dtype=torch.bfloat16), + 20, + 1e-6, + 1e-6, + 1e-6, + 2.0, + 1e-5, + ) + + +def _assert_like(actual: torch.Tensor, shape: tuple, dtype: torch.dtype) -> None: + assert tuple(actual.shape) == shape and actual.dtype == dtype, (actual.shape, actual.dtype) + assert actual.device.type == "cuda", actual.device + + +@pytest.mark.parametrize("rows", [1, decode_moe_module.MAX_ROWS]) +def test_decode_moe_under_fake_tensor_mode(rows): + with FakeTensorMode(): + inputs = _decode_moe_inputs(rows) + out = _empty(rows, HIDDEN, dtype=torch.bfloat16) + assert native_ops.decode_moe(*inputs, out, False) is None + wrapped = decode_moe_module.decode_moe(*inputs) + _assert_like(out, (rows, HIDDEN), torch.bfloat16) + _assert_like(wrapped, (rows, HIDDEN), torch.bfloat16) + + +@pytest.mark.parametrize("rows", [1, boundary_module.MAX_ROWS]) +def test_mhc_ar_boundary_under_fake_tensor_mode(rows): + outputs = ( + ((rows, MULT, HIDDEN), torch.bfloat16), + ((rows, MULT), torch.float32), + ((rows, MULT, MULT), torch.float32), + ((rows, HIDDEN), torch.bfloat16), + ) + with FakeTensorMode(): + inputs = _boundary_inputs(rows) + written = [_empty(*shape, dtype=dtype) for shape, dtype in outputs] + assert native_ops.mhc_ar_boundary(*inputs, *written, False) is None + wrapped = boundary_module.mhc_ar_boundary(*inputs) + assert len(wrapped) == len(outputs) + for (shape, dtype), out, returned in zip(outputs, written, wrapped): + _assert_like(out, shape, dtype) + _assert_like(returned, shape, dtype) diff --git a/tests/unittest/_torch/modeling_v2/models/glm5_next/test_glm5_next_routing.py b/tests/unittest/_torch/modeling_v2/models/glm5_next/test_glm5_next_routing.py index 9a6caacf3ba1..f4a5730e1713 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm5_next/test_glm5_next_routing.py +++ b/tests/unittest/_torch/modeling_v2/models/glm5_next/test_glm5_next_routing.py @@ -59,6 +59,22 @@ def _flash_config(**text_overrides) -> PretrainedConfig: n_routed_experts=288, q_lora_rank=1536, linear_num_heads=64, + num_experts_per_tok=8, + moe_intermediate_size=2048, + hc_mult=4, + scoring_func="sigmoid", + topk_method="noaux_tc", + n_group=1, + topk_group=1, + norm_topk_prob=True, + n_shared_experts=1, + tie_word_embeddings=False, + linear_attn_config={ + "num_heads": 64, + "head_dim": 128, + "short_conv_kernel_size": 4, + "gate_lower_bound": -5.0, + }, ) text_config.update(text_overrides) return PretrainedConfig( @@ -69,13 +85,17 @@ def _flash_config(**text_overrides) -> PretrainedConfig: def _model_config( - pretrained_config, quant_algo=QuantAlgo.NVFP4, spec_config=None, **mapping_kwargs + pretrained_config, + quant_algo=QuantAlgo.NVFP4, + spec_config=None, + kv_cache_quant_algo=None, + **mapping_kwargs, ): return ModelConfig( pretrained_config=pretrained_config, mapping=Mapping(**(mapping_kwargs or _TEP4)), spec_config=spec_config, - quant_config=QuantConfig(quant_algo=quant_algo), + quant_config=QuantConfig(quant_algo=quant_algo, kv_cache_quant_algo=kv_cache_quant_algo), ) @@ -177,3 +197,373 @@ def test_flash_on_sm103_does_not_match(monkeypatch): 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")) + + +# The block-FP8 target: the FP8_BLOCK_SCALES checkpoint on sm_100, tp4, FP8 KV cache, text only, +# with no speculation or one-model MTP-Eagle at three linear drafts. + +_FP8 = "models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.modeling.ModelingV2Glm53FlashFp8Sm10xTp4" +_TP4 = dict(world_size=4, tp_size=4, moe_tp_size=4, moe_ep_size=1, gpus_per_node=8) + + +def _fp8_config( + spec_config=None, kv_cache_quant_algo=QuantAlgo.FP8, text_overrides=None, **mapping +): + """The block-FP8 checkpoint at the deployment its target admits, unless overridden.""" + return _model_config( + _flash_config(**(text_overrides or {})), + quant_algo=QuantAlgo.FP8_BLOCK_SCALES, + spec_config=spec_config, + kv_cache_quant_algo=kv_cache_quant_algo, + **(mapping or _TP4), + ) + + +def _fp8_llm_args(mode: str, **kwargs) -> TorchLlmArgs: + kwargs.setdefault("disable_mm_encoder", True) + return _llm_args(mode, **kwargs) + + +def _mtp3(**fields) -> MTPDecodingConfig: + """MTP-Eagle at three drafts with ``fields`` set as given, bypassing the validators that would + reject some of them on their own (a draft checkpoint path that does not exist, a tree without + its sizes): routing reads the fields, not their consistency.""" + return MTPDecodingConfig(max_draft_len=3).model_copy(update=fields) + + +def _fp8_declined(config, llm_args_of, missed: str) -> str: + """``auto`` falls back to the built-in model; ``require`` raises with ``missed`` in the reason, + which it returns.""" + assert modeling_v2_resolve(config, llm_args_of("auto")) is None + with pytest.raises(ValueError, match=missed) as raised: + modeling_v2_resolve(config, llm_args_of("require")) + return str(raised.value) + + +@pytest.mark.parametrize("mode", ["auto", "require"]) +@pytest.mark.parametrize("spec_config", [None, _mtp3()], ids=["no-spec", "mtp3"]) +@pytest.mark.parametrize("moe_backend", ["AUTO", "TRTLLM"]) +@pytest.mark.parametrize("allreduce_strategy", ["AUTO", "ONESHOT"]) +def test_fp8_tp4_routes_to_the_fp8_target(allreduce_strategy, moe_backend, spec_config, mode): + llm_args = _fp8_llm_args( + mode, moe_config={"backend": moe_backend}, allreduce_strategy=allreduce_strategy + ) + assert _qualified(modeling_v2_resolve(_fp8_config(spec_config), llm_args)) == _FP8 + + +def test_fp8_kv_cache_dtype_override_routes_to_the_fp8_target(): + """Deployments ask for the FP8 latent cache with kv_cache_config.dtype; the checkpoint declares none.""" + llm_args = _fp8_llm_args("require", kv_cache_config={"dtype": "fp8"}) + config = _fp8_config(kv_cache_quant_algo=None) + assert _qualified(modeling_v2_resolve(config, llm_args)) == _FP8 + + +def test_fp8_target_is_built_only_when_modeling_v2_is_set(): + """The built-in model's defaults are applied before routing; they do not count as setting modeling_v2.""" + from tensorrt_llm._torch.models.modeling_glm5_next_vision import Glm5NextVLM + from tensorrt_llm.llmapi.llm_utils import apply_model_defaults_to_llm_args + + llm_args = TorchLlmArgs(model=_DUMMY_MODEL, disable_mm_encoder=True) + apply_model_defaults_to_llm_args(llm_args, Glm5NextVLM.get_model_defaults(llm_args)) + assert llm_args.allreduce_strategy == "ONESHOT" + assert modeling_v2_resolve(_fp8_config(_mtp3()), llm_args) is None + llm_args.modeling_v2 = "auto" + assert _qualified(modeling_v2_resolve(_fp8_config(_mtp3()), llm_args)) == _FP8 + + +@pytest.mark.parametrize( + "spec_config, missed", + [ + (_mtp3(use_dynamic_tree=True), "a dynamic draft tree"), + (_mtp3(max_total_draft_tokens=6), "draft trees"), + (_mtp3(speculative_model="/draft"), "a separate speculative checkpoint"), + (_mtp3(draft_lm_head_dtype="mxfp8"), r"draft_lm_head_dtype=mxfp8 \(only auto\)"), + (_mtp3(moe_backend="CUTLASS"), "draft MoE backend CUTLASS"), + (_mtp3(draft_len_schedule={8: 3, 32: 2}), "a draft-length schedule"), + (MTPDecodingConfig(max_draft_len=3, max_concurrency=16), "a draft-length schedule"), + ( + _mtp3(acceptance_rate_window_size=8, acceptance_rate_threshold=0.3), + "the acceptance-rate gate", + ), + (_mtp3(use_rejection_sampling=True), "rejection sampling"), + (_mtp3(use_relaxed_acceptance_for_thinking=True), "relaxed acceptance"), + ( + Eagle3DecodingConfig(max_draft_len=3, speculative_model="unused"), + r"Eagle3DecodingConfig \(only MTP\)", + ), + ( + NGramDecodingConfig(max_draft_len=3, max_matching_ngram_size=2), + r"NGramDecodingConfig \(only MTP\)", + ), + ], + ids=[ + "dynamic-tree", + "static-tree", + "draft-checkpoint", + "mxfp8-draft-head", + "draft-moe-backend", + "draft-length-schedule", + "max-concurrency", + "acceptance-rate-gate", + "rejection-sampling", + "relaxed-acceptance", + "eagle3", + "ngram", + ], +) +def test_fp8_drafter_options_it_does_not_build_are_declined(spec_config, missed): + _fp8_declined(_fp8_config(spec_config), _fp8_llm_args, missed) + + +@pytest.mark.parametrize( + "gate", + [dict(acceptance_rate_window_size=8), dict(acceptance_rate_threshold=0.3)], + ids=["window-only", "threshold-only"], +) +def test_fp8_half_set_acceptance_rate_gate_is_admitted(gate): + """The executor builds the gate only when both fields are set; one alone changes nothing.""" + llm_args = _fp8_llm_args("require") + assert _qualified(modeling_v2_resolve(_fp8_config(_mtp3(**gate)), llm_args)) == _FP8 + + +@pytest.mark.parametrize( + "moe_update, args_update, missed", + [ + (dict(backend="DEEPGEMM"), {}, r"MoE backend DEEPGEMM \(only AUTO or TRTLLM\)"), + (dict(backend="CUTLASS"), {}, r"MoE backend CUTLASS \(only AUTO or TRTLLM\)"), + (dict(load_balancer="eplb.yaml"), {}, r"a MoE load balancer \(eplb\.yaml\)"), + (dict(max_num_tokens=4096), {}, r"moe_config\.max_num_tokens=4096"), + ({}, dict(dwdp_config="dwdp"), r"DWDP \(dwdp\)"), + ], + ids=["deepgemm", "cutlass", "load-balancer", "moe-chunking", "dwdp"], +) +def test_fp8_moe_options_it_does_not_honour_are_declined(moe_update, args_update, missed): + def llm_args_of(mode): + args = _fp8_llm_args(mode) + moe_config = args.moe_config.model_copy(update=moe_update) + return args.model_copy(update=dict(args_update, moe_config=moe_config)) + + _fp8_declined(_fp8_config(), llm_args_of, missed) + + +def test_fp8_target_ignores_the_built_in_shared_expert_fusion_flag(monkeypatch): + """The flag configures the built-in MoE backends, which the target never builds.""" + monkeypatch.setenv("TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION", "1") + assert _qualified(modeling_v2_resolve(_fp8_config(), _fp8_llm_args("require"))) == _FP8 + + +@pytest.mark.parametrize( + "llm_kwargs, missed", + [ + (dict(prefill_cuda_graph_backend="breakable"), r"prefill CUDA graphs \(breakable\)"), + (dict(prefill_cuda_graph_backend="piecewise"), r"prefill CUDA graphs \(piecewise\)"), + (dict(torch_compile_config={}), "torch.compile"), + (dict(load_format="dummy"), r"load_format DUMMY \(only AUTO\)"), + (dict(lm_head_output_dtype="float32"), r"lm_head_output_dtype=float32 \(only auto\)"), + (dict(allreduce_strategy="NCCL"), r"allreduce_strategy=NCCL \(only AUTO or ONESHOT\)"), + (dict(kv_cache_config={"dtype": "fp8_ds_mla"}), r"KV cache fp8_ds_mla \(only FP8\)"), + ], + ids=[ + "breakable-prefill-graphs", + "piecewise-prefill-graphs", + "torch-compile", + "dummy-weights", + "fp32-logits", + "nccl-allreduce", + "fp8-ds-mla-cache", + ], +) +def test_fp8_engine_settings_it_does_not_run_are_declined(llm_kwargs, missed): + def llm_args_of(mode): + return _fp8_llm_args(mode, **llm_kwargs) + + _fp8_declined(_fp8_config(), llm_args_of, missed) + + +@pytest.mark.parametrize( + "config_kwargs, llm_kwargs, missed", + [ + ( + dict(text_overrides=dict(num_hidden_layers=78)), + {}, + r"text_config shape \(78, 4096, 288, 1536, 64\) \(only", + ), + (dict(kv_cache_quant_algo=None), {}, r"KV cache bf16 \(only FP8\)"), + (dict(kv_cache_quant_algo=QuantAlgo.NVFP4), {}, r"KV cache NVFP4 \(only FP8\)"), + (dict(world_size=4, tp_size=4, moe_ep_size=4, moe_tp_size=1), {}, "moe_tp=1 moe_ep=4"), + ( + dict(world_size=4, tp_size=4, moe_tp_size=4, enable_attention_dp=True), + {}, + "attention DP", + ), + (dict(world_size=8, tp_size=8, moe_tp_size=8), {}, "world_size=8 tp=8"), + ( + dict(world_size=4, tp_size=4, moe_tp_size=4, gpus_per_node=2), + {}, + r"gpus_per_node=2 \(only the 4 ranks of one node\)", + ), + ( + dict(text_overrides=dict(num_experts_per_tok=6)), + {}, + r"text_config geometry .*'num_experts_per_tok': 6", + ), + ( + dict(text_overrides=dict(moe_intermediate_size=1024)), + {}, + r"text_config geometry .*'moe_intermediate_size': 1024", + ), + (dict(text_overrides=dict(hc_mult=2)), {}, r"text_config geometry .*'hc_mult': 2"), + (dict(text_overrides=dict(scoring_func="softmax")), {}, "MoE routing .*softmax"), + ( + dict( + text_overrides=dict( + linear_attn_config={ + "num_heads": 64, + "head_dim": 128, + "short_conv_kernel_size": 4, + } + ) + ), + {}, + "a KDA gate without gate_lower_bound", + ), + ( + dict( + text_overrides=dict( + linear_attn_config={ + "num_heads": 64, + "head_dim": 64, + "short_conv_kernel_size": 4, + "gate_lower_bound": -5.0, + } + ) + ), + {}, + r"KDA .*'head_dim': 64", + ), + ( + dict( + text_overrides=dict( + linear_attn_config={ + "num_heads": 64, + "head_dim": 128, + "short_conv_kernel_size": 3, + "gate_lower_bound": -5.0, + } + ) + ), + {}, + r"KDA .*'short_conv_kernel_size': 3", + ), + ( + dict(text_overrides=dict(tie_word_embeddings=True)), + {}, + "an LM head tied to the embedding", + ), + (dict(spec_config=MTPDecodingConfig(max_draft_len=2)), {}, r"max_draft_len=2 \(only 3\)"), + ( + dict(spec_config=MTPDecodingConfig(max_draft_len=3, use_mtp_vanilla=True)), + {}, + "only one-model MTP-Eagle", + ), + ({}, dict(disable_mm_encoder=False), "the vision encoder"), + ], + ids=[ + "other-shape", + "bf16-kv", + "nvfp4-kv", + "expert-parallel", + "attention-dp", + "tp8", + "two-gpus-per-node", + "top-k", + "expert-intermediate", + "mhc-streams", + "softmax-routing", + "unbounded-gate", + "kda-head-dim", + "kda-conv-width", + "tied-embeddings", + "mtp2", + "vanilla-mtp", + "vision-tower", + ], +) +def test_fp8_near_misses_do_not_match(config_kwargs, llm_kwargs, missed): + def llm_args_of(mode): + return _fp8_llm_args(mode, **llm_kwargs) + + _fp8_declined(_fp8_config(**config_kwargs), llm_args_of, missed) + + +def test_fp8_refusal_names_only_what_failed(): + """A deployment that misses one requirement is told that one, not every fact it has.""" + reason = _fp8_declined( + _fp8_config(kv_cache_quant_algo=None), _fp8_llm_args, r"KV cache bf16 \(only FP8\)" + ) + for admitted in ("world_size", "moe_tp", "attention DP", "MoE backend", "vision", "draft"): + assert admitted not in reason, f"{admitted!r} is admitted but reported: {reason}" + + +@pytest.mark.parametrize("quant_algo", [QuantAlgo.NVFP4, None], ids=["nvfp4", "unquantized"]) +def test_other_checkpoints_do_not_reach_the_fp8_target(quant_algo): + """Another quantization of the same model never lands on the block-FP8 target.""" + config = _model_config( + _flash_config(), quant_algo=quant_algo, kv_cache_quant_algo=QuantAlgo.FP8, **_TP4 + ) + target = modeling_v2_resolve(config, _fp8_llm_args("auto")) + assert target is None or _qualified(target) != _FP8 + try: + target = modeling_v2_resolve(config, _fp8_llm_args("require")) + except ValueError: + return + assert _qualified(target) != _FP8 + + +@pytest.mark.parametrize("sm", [(9, 0), (10, 3), (12, 0)], ids=["sm90", "sm103", "sm120"]) +def test_fp8_on_other_gpus_does_not_match(monkeypatch, sm): + """The target has run on B200 only; a B300 falls back.""" + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: sm) + missed = rf"compute capability \({sm[0]}, {sm[1]}\) \(only sm_100\)" + _fp8_declined(_fp8_config(), _fp8_llm_args, missed) + + +def test_fp8_constructor_refuses_what_route_declines(): + """A direct construction raises on the facts ``route`` declines, before building anything.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.modeling import ( + ModelingV2Glm53FlashFp8Sm10xTp4, + ) + + config = _fp8_config(world_size=8, tp_size=8, moe_tp_size=8) + config.disable_mm_encoder = True + with pytest.raises(ValueError, match="world_size=8 tp=8"): + ModelingV2Glm53FlashFp8Sm10xTp4(config) + config = _fp8_config(_mtp3(use_rejection_sampling=True)) + with pytest.raises(ValueError, match="rejection sampling; the vision encoder"): + ModelingV2Glm53FlashFp8Sm10xTp4(config) + + +@pytest.mark.parametrize( + "multimodal_data, refused", + [ + (None, None), + ({}, None), + ({"image": {"pixel_values": torch.zeros(1)}}, "image"), + ({"video": {"pixel_values_videos": torch.zeros(1)}}, "video"), + ], + ids=["text", "empty-payload", "image", "video"], +) +def test_fp8_target_refuses_media_on_admission(multimodal_data, refused): + """The text-only target fails a media request on admission, as that request's own error.""" + from types import SimpleNamespace + + from tensorrt_llm._torch._experimental.modeling_v2.models.glm5_next.glm_5_3_flash_fp8__sm_10x__tp4.modeling import ( + ModelingV2Glm53FlashFp8Sm10xTp4, + ) + + request = SimpleNamespace(py_multimodal_data=multimodal_data) + if refused is None: + ModelingV2Glm53FlashFp8Sm10xTp4.validate_request(None, request) + return + with pytest.raises(ValueError, match=f"carries {refused} input"): + ModelingV2Glm53FlashFp8Sm10xTp4.validate_request(None, request) diff --git a/tests/unittest/_torch/pyexecutor/engine/runners/decoder/test_runner.py b/tests/unittest/_torch/pyexecutor/engine/runners/decoder/test_runner.py index fc5cbb092b04..9d022fb58395 100644 --- a/tests/unittest/_torch/pyexecutor/engine/runners/decoder/test_runner.py +++ b/tests/unittest/_torch/pyexecutor/engine/runners/decoder/test_runner.py @@ -567,6 +567,8 @@ def test_global_incompatibilities_bypass_candidate_selection(self) -> None: "extra_inputs", "ple_recurrent_state", "nested_ple_recurrent_state", + "declared_recurrent_state", + "declared_recurrent_state_on_the_wrapper", "context_parallel", "context_only_cache", ) @@ -597,6 +599,12 @@ def test_global_incompatibilities_bypass_candidate_selection(self) -> None: engine._model.model = SimpleNamespace( llm=SimpleNamespace(model=SimpleNamespace(has_ple=True)) ) + elif case == "declared_recurrent_state": + engine._model.resets_recurrent_state_in_context = True + elif case == "declared_recurrent_state_on_the_wrapper": + # A causal-LM wrapper declares it; the decoder inside it does not. + engine._model.model = torch.nn.Identity() + engine._model.resets_recurrent_state_in_context = True elif case == "context_parallel": engine._mapping.cp_size = 2 elif case == "context_only_cache": diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 1e4256b8f3d6..8e9d4844ea0f 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -1118,6 +1118,18 @@ def test_compute_applied_llm_defaults_simple_field(self): applied = apply_model_defaults_to_llm_args(llm_args, model_defaults) assert applied == model_defaults + def test_applied_defaults_leave_other_fields_unset(self): + llm_args = TorchLlmArgs(model="/tmp/dummy_model", max_batch_size=8) + before = set(llm_args.model_fields_set) + apply_model_defaults_to_llm_args(llm_args, { + "enable_chunked_prefill": False, + "max_batch_size": 4, + }) + assert llm_args.max_batch_size == 8 + assert not llm_args.enable_chunked_prefill + assert llm_args.model_fields_set == before | {"enable_chunked_prefill"} + assert "modeling_v2" not in llm_args.model_fields_set + @pytest.mark.parametrize( "defaults_dict,should_raise,error_contains", [