diff --git a/3rdparty/fetch_content.json b/3rdparty/fetch_content.json index 24c079b57a5a..b9d3f9c4cb06 100644 --- a/3rdparty/fetch_content.json +++ b/3rdparty/fetch_content.json @@ -46,7 +46,10 @@ "git_tag": "b64107f2b9599ca76445b7f62eedf66bae1d095b", "git_submodules_recurse": true, "source_subdir": "dont-add-this-project-with-add-subdirectory", - "patch_file": "patches/deepgemm_configurable_barrier_timeout.patch" + "patch_file": [ + "patches/deepgemm_configurable_barrier_timeout.patch", + "patches/deepgemm_sm10x_family_target.patch" + ] }, { "name": "eigen", diff --git a/3rdparty/patches/deepgemm_sm10x_family_target.patch b/3rdparty/patches/deepgemm_sm10x_family_target.patch new file mode 100644 index 000000000000..44b9282fe76b --- /dev/null +++ b/3rdparty/patches/deepgemm_sm10x_family_target.patch @@ -0,0 +1,13 @@ +--- a/csrc/runtime/jit.hpp ++++ b/csrc/runtime/jit.hpp +@@ -37,6 +37,10 @@ + // Preserve the existing 12.x family target when DeepJIT derives a different minor target. + if (runtime->device.get_arch_major() == 12 and runtime->device.get_arch_minor() != 0) + runtime->default_compiler_options.arch = "120f"; ++ // A 10.x device past 10.3 (e.g. 10.7) builds for the 10.0 family target, which it runs: ++ // the bundled CUTLASS enables TMA and tcgen05 only for the targets it knows. ++ if (runtime->device.get_arch_major() == 10 and runtime->device.get_arch_minor() > 3) ++ runtime->default_compiler_options.arch = "100f"; + return runtime; + }); + } diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cu b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cu index af13ae72a61c..6545a3a4da23 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cu +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cu @@ -184,19 +184,20 @@ CUresult copyDeviceToDevice(std::vector const& tasks, ssize_t numBytes, return launchBatchedCopy(false, tasks, numBytes, stream); } -// dst_tensor[:, :num_seqs, 0] = src_tensor[:, copy_idx] -// dst_tensor[:, :num_seqs, 1] = dst_tensor[:, :num_seqs, 0] + 1 +// dst_tensor[:, :num_seqs, 0, :numBlocksToCopy] = src_tensor[:, copy_idx, 0, :numBlocksToCopy] +// dst_tensor[:, :num_seqs, 1, :numBlocksToCopy] = dst_tensor[:, :num_seqs, 0, :numBlocksToCopy] + 1 +// Rows are numBlocksPerSeq long in both tensors; the columns past numBlocksToCopy are not written. __global__ void copyBatchBlockOffsetsToDeviceKernel(SizeType32 const* __restrict__ srcPtr, SizeType32* __restrict__ dstPtr, SizeType32 const srcMaxNumSequences, SizeType32 const dstMaxNumSequences, - SizeType32 numBlocksPerSeq, SizeType32 const* __restrict__ copyIndex, SizeType32 const* __restrict__ indexScales, - SizeType32 const* __restrict__ kvOffset) + SizeType32 numBlocksPerSeq, SizeType32 numBlocksToCopy, SizeType32 const* __restrict__ copyIndex, + SizeType32 const* __restrict__ indexScales, SizeType32 const* __restrict__ kvOffset) { constexpr uint32_t kvFactor = 2; constexpr auto elemPerAccess = sizeof(PackedInt) / sizeof(SizeType32); __shared__ PackedInt data[copyBlocknbBufs][copyBlockCtaSize]; - auto const iterPerSeq = divUp(numBlocksPerSeq * sizeof(SizeType32), sizeof(PackedInt) * copyBlockCtaSize); + auto const iterPerSeq = divUp(numBlocksToCopy * sizeof(SizeType32), sizeof(PackedInt) * copyBlockCtaSize); auto const tid = threadIdx.x; auto const poolIdx = blockIdx.x; auto const seqIdx = blockIdx.y; @@ -205,7 +206,7 @@ __global__ void copyBatchBlockOffsetsToDeviceKernel(SizeType32 const* __restrict uint32_t const dstIdxKBeg = tid * elemPerAccess + (poolIdx * dstMaxNumSequences + seqIdx) * seqDimStride; uint32_t const dstIdxVBeg = dstIdxKBeg + numBlocksPerSeq; - uint32_t const srcIdxEnd = (poolIdx * srcMaxNumSequences + copyIndex[seqIdx]) * seqDimStride + numBlocksPerSeq; + uint32_t const srcIdxEnd = (poolIdx * srcMaxNumSequences + copyIndex[seqIdx]) * seqDimStride + numBlocksToCopy; for (uint32_t i = 0; i < iterPerSeq + copyBlocknbBufs; i++) { @@ -250,7 +251,7 @@ __global__ void copyBatchBlockOffsetsToDeviceKernel(SizeType32 const* __restrict // Host-side launcher void copyBatchBlockOffsetsToDevice(ITensor const& input, ITensor& output, ITensor const& copyIndex, - ITensor const& indexScales, ITensor const& kvOffset, CUstream stream) noexcept + ITensor const& indexScales, ITensor const& kvOffset, SizeType32 maxNumBlocks, CUstream stream) noexcept { using namespace tensorrt_llm::runtime; @@ -282,11 +283,22 @@ void copyBatchBlockOffsetsToDevice(ITensor const& input, ITensor& output, ITenso "Not implemented case: numBlocksPerSeq * sizeof(SizeType32) = %zu must be a multiple of %zu.", static_cast(numBlocksPerSeq * sizeof(SizeType32)), static_cast(sizeof(PackedInt))); + // A negative maxNumBlocks copies whole rows. Otherwise the copy covers the first maxNumBlocks + // columns, rounded up to whole PackedInt accesses, which stay inside the row. + SizeType32 numBlocksToCopy = numBlocksPerSeq; + if (maxNumBlocks >= 0) + { + constexpr SizeType32 blocksPerAccess = sizeof(PackedInt) / sizeof(SizeType32); + SizeType32 const roundedBlocks + = (std::max(maxNumBlocks, SizeType32{1}) + blocksPerAccess - 1) / blocksPerAccess * blocksPerAccess; + numBlocksToCopy = std::min(numBlocksPerSeq, roundedBlocks); + } + dim3 gridDim(numPools, numSeqs, 1); dim3 blockDim(copyBlockCtaSize); copyBatchBlockOffsetsToDeviceKernel<<>>(srcPtr, dstPtr, srcMaxNumSequences, - dstMaxNumSequences, numBlocksPerSeq, copyIndexPtr, indexScalesPtr, kvOffsetPtr); + dstMaxNumSequences, numBlocksPerSeq, numBlocksToCopy, copyIndexPtr, indexScalesPtr, kvOffsetPtr); } } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h index f15e96daf10b..e8f4308f120f 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h @@ -108,7 +108,9 @@ CUresult copyHostToDevice(std::vector> const& tasks CUresult copyDeviceToHost(std::vector> const& tasks, ssize_t numBytes, CUstream stream); CUresult copyDeviceToDevice(std::vector> const& tasks, ssize_t numBytes, CUstream stream); +//! Copies the block offsets of the batch rows. A non-negative maxNumBlocks bounds the copied columns +//! (rounded up to a multiple of four); a negative one copies whole rows. void copyBatchBlockOffsetsToDevice(ITensor const& input, ITensor& output, ITensor const& copyIndex, - ITensor const& indexScales, ITensor const& kvOffset, CUstream stream) noexcept; + ITensor const& indexScales, ITensor const& kvOffset, SizeType32 maxNumBlocks, CUstream stream) noexcept; } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.cu new file mode 100644 index 000000000000..2a5bb29cdad2 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.cu @@ -0,0 +1,607 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Weight-streaming BF16 kernels for the dense (non-MoE) MLP at decode row counts (M <= 8): the gate/up projection +// with SwiGLU, and the down projection. The compiled tilings serve hidden 6144 with a rank-local intermediate size of +// 768, the dense layers 0-2 of GLM-5.2 at TP16; the result is the rank-local partial that the post-MLP all-reduce sums. +// +// Each CTA owns 16-row weight tiles and its warps split the contraction. For every 32-wide contraction step a thread +// takes 16 bytes of two weight rows per tile and of one activation row and feeds them to two m16n8k16 MMAs, the +// weight on the MMA M side and the (at most 8) activation rows on the N side. The contraction index of an MMA k-slot +// depends only on the thread's position inside its quad, so the weight and activation fragments use the same +// permutation of k: the result is the same sum in another FP32 order. The gate/up kernel holds the gate rows and the +// up rows of the same output columns, and the thread that finishes a column applies SwiGLU in registers. +// +// Weights never depend on the preceding kernel: before griddepcontrol.wait each CTA copies all of its weights into +// shared memory with one bulk copy per weight row. After the wait each thread loads its activations from L2 and runs +// its MMAs; the warp sums are added in warp order in shared memory. With kClusterK > 1 the CTAs of a cluster split +// the contraction: each sends its sums to the first CTA of the cluster with st.async, which adds them in rank order. +// The results do not depend on timing. +// +// Programmatic dependent launch: activations are read only after griddepcontrol.wait; the dependents are released at +// entry or after the weight copies are issued (Trigger). Each grid fits in one wave of resident CTAs. + +#include "tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::dense_mlp_small_m +{ + +namespace +{ + +using Bf16 = __nv_bfloat16; + +constexpr int kWarpSize = 32; +//! Weight rows of one MMA tile (MMA M). +constexpr int kTileRows = 16; +//! BF16 values per 16-byte vector load. +constexpr int kVecElems = 8; +//! Threads per MMA quad; a quad covers kQuadThreads * kVecElems contraction elements per step. +constexpr int kQuadThreads = 4; +//! Contraction elements consumed by one warp step (two m16n8k16 MMAs per row tile). +constexpr int kKPerStep = kQuadThreads * kVecElems; +constexpr int kStepBytes = kKPerStep * static_cast(sizeof(Bf16)); + +//! When the dependents of a grid are released. +enum class Trigger : int +{ + kAtEntry, //!< first instruction + kAfterIssue, //!< after the weight copies are issued, before griddepcontrol.wait +}; + +//! Compile-time tiling of one kernel instance. +//! \tparam kGateUp gate/up + SwiGLU (true) or down projection (false). +//! \tparam kRowTiles 16-row weight tiles per CTA. gate/up: one tile holds 8 gate and the matching 8 up rows; with +//! more tiles, the first half are gate rows and the second half the up rows of the same columns. +//! \tparam kWarps warps per CTA; each warp owns a contiguous slice of the CTA's contraction. +//! \tparam kSteps 32-wide contraction steps per warp. +//! \tparam kClusterK CTAs of a cluster that split the contraction. +template +struct Tiling +{ + static constexpr bool kGateUp = kGateUp_; + static constexpr int kRowTiles = kRowTiles_; + static constexpr int kWarps = kWarps_; + static constexpr int kSteps = kSteps_; + static constexpr int kClusterK = kClusterK_; + static constexpr Trigger kTrigger = kTrigger_; + + static constexpr int kThreads = kWarps * kWarpSize; + static constexpr int kRows = kRowTiles * kTileRows; + static constexpr int kKPerWarp = kSteps * kKPerStep; + static constexpr int kKPerCta = kWarps * kKPerWarp; + //! The contraction size this tiling serves. + static constexpr int kKDim = kKPerCta * kClusterK; + //! Output columns of a CTA (gate/up: of h, down: of out). + static constexpr int kOutCols = kGateUp ? kRows / 2 : kRows; + //! (column, token) values the cluster finishes; each holds kEpiVals sums (gate and up, or one). + static constexpr int kEpi = kOutCols * kMaxRows; + static constexpr int kEpiVals = kGateUp ? 2 : 1; + + //! Weight stage [row][kKPerCta]: the bytes of one row copy, rows padded to 64 mod 128 bytes so that the 16-byte + //! fragment loads of a quarter warp (two rows, four quads) hit distinct banks. + static constexpr int kRowBytes = kKPerCta * static_cast(sizeof(Bf16)); + static constexpr int kRowStride = kRowBytes % 128 == 0 ? kRowBytes + 64 : kRowBytes; + + // Shared-memory layout (bytes): warp sums | sums received by the cluster leader | weight stage | barriers. + static constexpr int kRedBytes = kWarps * kRowTiles * 4 * kWarpSize * static_cast(sizeof(float)); + static constexpr int kRecvBytes = (kClusterK - 1) * kEpi * kEpiVals * static_cast(sizeof(float)); + static constexpr int kStageOffset = (kRedBytes + kRecvBytes + 127) / 128 * 128; + static constexpr int kBarOffset = kStageOffset + kRows * kRowStride; + static constexpr int kSmemBytes = kBarOffset + 2 * static_cast(sizeof(uint64_t)); + + static_assert(!kGateUp_ || kRowTiles_ == 1 || kRowTiles_ % 2 == 0, "gate/up needs one or an even number of tiles"); + static_assert(kRows <= kThreads, "one thread issues the copy of each weight row"); + static_assert(kRowBytes % 64 == 0, "rows are whole contraction steps"); +}; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + +__device__ __forceinline__ void pdlWait() +{ + asm volatile("griddepcontrol.wait;" ::: "memory"); +} + +__device__ __forceinline__ void pdlLaunchDependents() +{ + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +} + +//! Activation load from L2 (the preceding kernel wrote it; L1 is bypassed). +__device__ __forceinline__ uint4 loadActivation(Bf16 const* ptr) +{ + uint4 r; + asm volatile("ld.global.cg.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w) + : "l"(ptr) + : "memory"); + return r; +} + +__device__ __forceinline__ uint32_t smemAddress(void const* ptr) +{ + return static_cast(__cvta_generic_to_shared(ptr)); +} + +//! One arrival (this one) plus expectBytes of transactions. +__device__ __forceinline__ void mbarrierInit(uint32_t bar, uint32_t expectBytes) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(bar) : "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(bar), "r"(expectBytes) : "memory"); +} + +__device__ __forceinline__ void mbarrierWait(uint32_t bar) +{ + asm volatile( + "{\n" + ".reg .pred done;\n" + "WAIT_%=:\n" + "mbarrier.try_wait.parity.shared::cta.b64 done, [%0], 0;\n" + "@!done bra WAIT_%=;\n" + "}\n" ::"r"(bar) + : "memory"); +} + +__device__ __forceinline__ void mbarrierWaitCluster(uint32_t bar) +{ + asm volatile( + "{\n" + ".reg .pred done;\n" + "WAIT_%=:\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 done, [%0], 0;\n" + "@!done bra WAIT_%=;\n" + "}\n" ::"r"(bar) + : "memory"); +} + +__device__ __forceinline__ void bulkCopy(uint32_t dst, void const* src, uint32_t bytes, uint32_t bar) +{ + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"(dst), + "l"(src), "r"(bytes), "r"(bar) + : "memory"); +} + +__device__ __forceinline__ uint32_t clusterCtaRank() +{ + uint32_t rank; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(rank)); + return rank; +} + +__device__ __forceinline__ uint32_t mapToCta(uint32_t addr, uint32_t rank) +{ + uint32_t peer; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(peer) : "r"(addr), "r"(rank)); + return peer; +} + +__device__ __forceinline__ void mmaBf16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) +{ + asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +//! Two MMAs of one step and row tile: MMA A takes contraction elements 0..3 of the thread's 8-element vectors, MMA B +//! elements 4..7. Fragments: a0/a2 = weight row "group", a1/a3 = weight row "group + 8", b0/b1 = activation row +//! "group". +__device__ __forceinline__ void mmaStep( + float (&accA)[4], float (&accB)[4], uint4 const& wLo, uint4 const& wHi, uint4 const& xv) +{ + mmaBf16(accA, wLo.x, wHi.x, wLo.y, wHi.y, xv.x, xv.y); + mmaBf16(accB, wLo.z, wHi.z, wLo.w, wHi.w, xv.z, xv.w); +} + +__device__ __forceinline__ float roundToBf16(float v) +{ + return __bfloat162float(__float2bfloat16_rn(v)); +} + +//! SwiGLU of BF16-rounded gate and up sums, as the BF16 GEMM output feeding silu_and_mul: the sigmoid takes the fast +//! exp and division of that kernel. +__device__ __forceinline__ float swiglu(float gateSum, float upSum) +{ + float const gate = roundToBf16(gateSum); + float const up = roundToBf16(upSum); + float const sig = __fdividef(1.F, 1.F + __expf(-gate)); + return gate * sig * up; +} + +//! Weight row (of the [*, kDim] weight matrix) of row r of row tile t of the CTA of output tile `tile`. +template +__device__ __forceinline__ int weightRow(int tile, int t, int r, int inter) +{ + if constexpr (T::kGateUp) + { + if constexpr (T::kRowTiles == 1) + { + int const col = tile * T::kOutCols + (r & 7); + return r < 8 ? col : inter + col; + } + else + { + constexpr int kHalf = T::kRowTiles / 2; + int const col = tile * T::kOutCols + (t % kHalf) * kTileRows + r; + return t < kHalf ? col : inter + col; + } + } + else + { + return tile * T::kRows + t * kTileRows + r; + } +} + +#endif // __CUDA_ARCH__ >= 900 + +//! x [numRows, kDim], weight [*, kDim]; out: gate/up h [numRows, outDim] (outDim = inter), down out [numRows, outDim]. +template +__global__ void __launch_bounds__(T::kThreads, 1) denseMlpSmallMKernel( + Bf16 const* __restrict__ x, Bf16 const* __restrict__ weight, Bf16* __restrict__ out, int numRows, int outDim) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + constexpr int kDim = T::kKDim; + extern __shared__ __align__(128) char smem[]; + float* red = reinterpret_cast(smem); + [[maybe_unused]] float* recv = reinterpret_cast(smem + T::kRedBytes); + char* stage = smem + T::kStageOffset; + uint32_t const stageBar = smemAddress(smem + T::kBarOffset); + [[maybe_unused]] uint32_t const clusterBar = stageBar + static_cast(sizeof(uint64_t)); + + if constexpr (T::kTrigger == Trigger::kAtEntry) + { + pdlLaunchDependents(); + } + + int const lane = threadIdx.x % kWarpSize; + int const warp = threadIdx.x / kWarpSize; + int const group = lane / kQuadThreads; + int const tig = lane % kQuadThreads; + int const rank = T::kClusterK > 1 ? static_cast(clusterCtaRank()) : 0; + int const tile = blockIdx.x / T::kClusterK; + int const kCta = rank * T::kKPerCta; + + // The weight stage expects all row bytes; the cluster leader also expects the sums of the other CTAs. + if (threadIdx.x == 0) + { + mbarrierInit(stageBar, static_cast(T::kRows * T::kRowBytes)); + if constexpr (T::kClusterK > 1) + { + if (rank == 0) + { + mbarrierInit(clusterBar, static_cast(T::kRecvBytes)); + } + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + __syncthreads(); + if constexpr (T::kClusterK > 1) + { + // Publishes the leader's barrier; every CTA waits for this phase before it sends. + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); + } + + // Thread r copies weight row r of the CTA (its contraction slice). + int const copyRow = static_cast(threadIdx.x); + if (copyRow < T::kRows) + { + int const wRow = weightRow(tile, copyRow / kTileRows, copyRow % kTileRows, outDim); + bulkCopy(smemAddress(stage + copyRow * T::kRowStride), weight + static_cast(wRow) * kDim + kCta, + static_cast(T::kRowBytes), stageBar); + } + + if constexpr (T::kTrigger == Trigger::kAfterIssue) + { + pdlLaunchDependents(); + } + pdlWait(); + + // Activations: row "group" (zero beyond numRows), the thread's 8 contraction elements of every step. + int const kWarp = kCta + warp * T::kKPerWarp; + bool const hasRow = group < numRows; + Bf16 const* xRow = x + static_cast(hasRow ? group : 0) * kDim + kWarp + tig * kVecElems; + uint4 xv[T::kSteps]; +#pragma unroll + for (int s = 0; s < T::kSteps; ++s) + { + xv[s] = hasRow ? loadActivation(xRow + s * kKPerStep) : make_uint4(0U, 0U, 0U, 0U); + } + + float accA[T::kRowTiles][4]; + float accB[T::kRowTiles][4]; +#pragma unroll + for (int t = 0; t < T::kRowTiles; ++t) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + accA[t][e] = 0.F; + accB[t][e] = 0.F; + } + } + mbarrierWait(stageBar); + char const* seg = stage + warp * T::kKPerWarp * static_cast(sizeof(Bf16)) + tig * 16; +#pragma unroll + for (int s = 0; s < T::kSteps; ++s) + { +#pragma unroll + for (int t = 0; t < T::kRowTiles; ++t) + { + int const rowLo = t * kTileRows + group; + uint4 const wLo = *reinterpret_cast(seg + rowLo * T::kRowStride + s * kStepBytes); + uint4 const wHi = *reinterpret_cast(seg + (rowLo + 8) * T::kRowStride + s * kStepBytes); + mmaStep(accA[t], accB[t], wLo, wHi, xv[s]); + } + } + + // Warp sums -> shared memory [warp][tile * 4 + element][lane]. +#pragma unroll + for (int t = 0; t < T::kRowTiles; ++t) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + red[(warp * T::kRowTiles * 4 + t * 4 + e) * kWarpSize + lane] = accA[t][e] + accB[t][e]; + } + } + __syncthreads(); + if constexpr (T::kClusterK > 1) + { + asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); + } + + // Finish value j: element (tile t, accumulator e, lane l) holds weight row t * 16 + l / 4 + (e / 2) * 8 of the + // CTA and token 2 * (l % 4) + e % 2. + for (int j = static_cast(threadIdx.x); j < T::kEpi; j += T::kThreads) + { + int q0; + [[maybe_unused]] int q1 = 0; + int col; + int token; + int const l = j % kWarpSize; + if constexpr (T::kGateUp && T::kRowTiles == 1) + { + int const e = j / kWarpSize; // 0 or 1: rows l / 4 (gate) and l / 4 + 8 (up) + q0 = e; + q1 = e + 2; + col = l / kQuadThreads; + token = 2 * (l % kQuadThreads) + e; + } + else if constexpr (T::kGateUp) + { + int const t = j / (4 * kWarpSize); + int const e = (j / kWarpSize) % 4; + q0 = t * 4 + e; + q1 = (t + T::kRowTiles / 2) * 4 + e; + col = t * kTileRows + l / kQuadThreads + (e / 2) * 8; + token = 2 * (l % kQuadThreads) + e % 2; + } + else + { + int const q = j / kWarpSize; + q0 = q; + col = (q / 4) * kTileRows + l / kQuadThreads + ((q % 4) / 2) * 8; + token = 2 * (l % kQuadThreads) + q % 2; + } + // Fixed order over warps keeps the result deterministic. + float v0 = 0.F; + [[maybe_unused]] float v1 = 0.F; +#pragma unroll + for (int w = 0; w < T::kWarps; ++w) + { + v0 += red[(w * T::kRowTiles * 4 + q0) * kWarpSize + l]; + if constexpr (T::kGateUp) + { + v1 += red[(w * T::kRowTiles * 4 + q1) * kWarpSize + l]; + } + } + if constexpr (T::kClusterK > 1) + { + if (rank != 0) + { + uint32_t const slot = smemAddress(recv + ((rank - 1) * T::kEpi + j) * T::kEpiVals); + uint32_t const peerSlot = mapToCta(slot, 0); + uint32_t const peerBar = mapToCta(clusterBar, 0); + if constexpr (T::kGateUp) + { + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.b32 [%0], {%1, %2}, [%3];" ::"r"( + peerSlot), + "r"(__float_as_uint(v0)), "r"(__float_as_uint(v1)), "r"(peerBar) + : "memory"); + } + else + { + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(peerSlot), + "r"(__float_as_uint(v0)), "r"(peerBar) + : "memory"); + } + continue; + } + mbarrierWaitCluster(clusterBar); + // Rank order. +#pragma unroll + for (int r = 1; r < T::kClusterK; ++r) + { + v0 += recv[((r - 1) * T::kEpi + j) * T::kEpiVals]; + if constexpr (T::kGateUp) + { + v1 += recv[((r - 1) * T::kEpi + j) * T::kEpiVals + 1]; + } + } + } + if (token < numRows) + { + int64_t const idx = static_cast(token) * outDim + tile * T::kOutCols + col; + if constexpr (T::kGateUp) + { + out[idx] = __float2bfloat16_rn(swiglu(v0, v1)); + } + else + { + out[idx] = __float2bfloat16_rn(v0); + } + } + } +#endif +} + +template +cudaError_t launchTiling( + Bf16 const* in, Bf16 const* weight, Bf16* out, int numRows, int outDim, bool enablePdl, cudaStream_t stream) +{ + auto const kernel = &denseMlpSmallMKernel; + if constexpr (T::kSmemBytes > 48 * 1024) + { + static cudaError_t const attrStatus + = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, T::kSmemBytes); + if (attrStatus != cudaSuccess) + { + return attrStatus; + } + } + cudaLaunchConfig_t config{}; + config.gridDim = dim3(outDim / T::kOutCols * T::kClusterK); + config.blockDim = dim3(T::kThreads); + config.dynamicSmemBytes = T::kSmemBytes; + config.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enablePdl ? 1 : 0; + attrs[1].id = cudaLaunchAttributeClusterDimension; + attrs[1].val.clusterDim.x = T::kClusterK; + attrs[1].val.clusterDim.y = 1; + attrs[1].val.clusterDim.z = 1; + config.attrs = attrs; + config.numAttrs = T::kClusterK > 1 ? 2 : 1; + return cudaLaunchKernelEx(&config, kernel, in, weight, out, numRows, outDim); +} + +// Tilings: row tiles, warps, steps, cluster size, trigger. +using GateUpTilings = std::tuple< + // Hidden 6144 (GLM-5.2): 2-CTA clusters split the contraction, 16 warps x 6 steps, 8 columns per cluster; at TP16 + // (inter 768) 192 CTAs with 96 KiB of weights each. The down kernel is released at entry, so its weights are staged + // during the preceding all-reduce as well. + Tiling>; + +using DownTilings = std::tuple< + // Inter 768 (GLM-5.2 at TP16): 64 rows per CTA, 8 warps x 3 steps; at hidden 6144 96 CTAs with 96 KiB of weights + // each. + Tiling>; + +template +constexpr int kNumTilings = static_cast(std::tuple_size_v); + +template +bool tilingSupported(int tiling, int numRows, int kDim, int outDim) +{ + if constexpr (I < kNumTilings) + { + if (tiling == I) + { + using T = std::tuple_element_t; + return numRows >= 1 && numRows <= kMaxRows && kDim == T::kKDim && outDim > 0 && outDim % T::kOutCols == 0; + } + return tilingSupported(tiling, numRows, kDim, outDim); + } + else + { + return false; + } +} + +template +cudaError_t launchTilingIndex(int tiling, Bf16 const* in, Bf16 const* weight, Bf16* out, int numRows, int outDim, + bool enablePdl, cudaStream_t stream) +{ + if constexpr (I < kNumTilings) + { + if (tiling == I) + { + return launchTiling>(in, weight, out, numRows, outDim, enablePdl, stream); + } + return launchTilingIndex(tiling, in, weight, out, numRows, outDim, enablePdl, stream); + } + else + { + return cudaErrorInvalidValue; + } +} + +//! The tiling a call runs: the requested one if it serves the shape, for kDefaultTiling the first entry that does; +//! -1 if none. +template +int resolveTiling(int tiling, int numRows, int kDim, int outDim) +{ + if (tiling != kDefaultTiling) + { + return tilingSupported(tiling, numRows, kDim, outDim) ? tiling : -1; + } + for (int i = 0; i < kNumTilings; ++i) + { + if (tilingSupported(i, numRows, kDim, outDim)) + { + return i; + } + } + return -1; +} + +} // namespace + +bool isGateUpSupported(int numRows, int hidden, int inter, int tiling) +{ + return resolveTiling(tiling, numRows, hidden, inter) >= 0; +} + +bool isDownSupported(int numRows, int inter, int hidden, int tiling) +{ + return resolveTiling(tiling, numRows, inter, hidden) >= 0; +} + +cudaError_t launchGateUpSwiglu(Bf16 const* x, Bf16 const* gateUpWeight, Bf16* h, int numRows, int hidden, int inter, + int tiling, bool enablePdl, cudaStream_t stream) +{ + int const resolved = resolveTiling(tiling, numRows, hidden, inter); + if (resolved < 0) + { + return cudaErrorInvalidValue; + } + return launchTilingIndex(resolved, x, gateUpWeight, h, numRows, inter, enablePdl, stream); +} + +cudaError_t launchDown(Bf16 const* h, Bf16 const* downWeight, Bf16* out, int numRows, int inter, int hidden, int tiling, + bool enablePdl, cudaStream_t stream) +{ + int const resolved = resolveTiling(tiling, numRows, inter, hidden); + if (resolved < 0) + { + return cudaErrorInvalidValue; + } + return launchTilingIndex(resolved, h, downWeight, out, numRows, hidden, enablePdl, stream); +} + +} // namespace kernels::glm_moe_dsa::dense_mlp_small_m + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.h new file mode 100644 index 000000000000..968c74dc1b55 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.h @@ -0,0 +1,62 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::dense_mlp_small_m +{ + +//! Largest number of activation rows (tokens) the kernels accept. +constexpr int kMaxRows = 8; + +//! Tiling argument that selects the default tiling of the shape. +constexpr int kDefaultTiling = -1; + +//! Host-side eligibility checks (no device work). tiling = kDefaultTiling checks the default tiling of the shape. +//! \param numRows activation rows M. +//! \param hidden model hidden size (gate/up contraction size, down output size). +//! \param inter rank-local intermediate size (gate/up output size, down contraction size). +bool isGateUpSupported(int numRows, int hidden, int inter, int tiling = kDefaultTiling); +bool isDownSupported(int numRows, int inter, int hidden, int tiling = kDefaultTiling); + +//! Gate/up projection fused with SwiGLU. For every row m < numRows and column j < inter: +//! gate = bf16(sum_k x[m, k] * W[j, k]), up = bf16(sum_k x[m, k] * W[inter + j, k]), +//! h[m, j] = bf16((gate * sigmoid(gate)) * up), sums in FP32, sigmoid as 1 / (1 + exp(-gate)) with the fast exp +//! and division. +//! \param x activations [numRows, hidden], row-major, BF16, 16-byte aligned. +//! \param gateUpWeight [2 * inter, hidden], row-major, BF16 (rows [0, inter) gate, [inter, 2 * inter) up). +//! \param h output [numRows, inter], row-major, BF16. +cudaError_t launchGateUpSwiglu(__nv_bfloat16 const* x, __nv_bfloat16 const* gateUpWeight, __nv_bfloat16* h, int numRows, + int hidden, int inter, int tiling, bool enablePdl, cudaStream_t stream); + +//! Down projection: out[m, n] = bf16(sum_k h[m, k] * W[n, k]) with FP32 sums. +//! \param h activations [numRows, inter], row-major, BF16, 16-byte aligned. +//! \param downWeight [hidden, inter], row-major, BF16. +//! \param out output [numRows, hidden], row-major, BF16. +cudaError_t launchDown(__nv_bfloat16 const* h, __nv_bfloat16 const* downWeight, __nv_bfloat16* out, int numRows, + int inter, int hidden, int tiling, bool enablePdl, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa::dense_mlp_small_m + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.cu new file mode 100644 index 000000000000..5a936de3083c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.cu @@ -0,0 +1,1000 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/common/lamportUtils.cuh" + +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ +namespace +{ + +using tensorrt_llm::common::loadPackedVolatile; +using tensorrt_llm::common::VolatilePackedLoad; + +constexpr int kBlockSize = 512; +constexpr int kNumWarps = kBlockSize / 32; +// Logit vectors (8 BF16) each thread loads before it compares them. +constexpr int kScanBatch = 8; +// FlashInfer RMSNormKernel geometry for H = 6144: thread t of a 64-thread group owns the 8-element vectors at columns +// 8 t + 512 b, b = 0..11. +constexpr int kNormThreads = 64; +constexpr int kNormVecElts = 8; +constexpr int kNormVecs = kDraftSampleGlueHidden / (kNormThreads * kNormVecElts); +constexpr int kNumNormGroups = kBlockSize / kNormThreads; +static_assert(kNormVecs * kNormThreads * kNormVecElts == kDraftSampleGlueHidden, "hidden size must be 512-aligned"); +// An unwritten word of the Lamport workspace (the MNNVL all-gather's empty word). +constexpr uint32_t kEmptyWord = 0xFFFFFFFFU; +constexpr int kMaxSlots = kDraftSampleGlueMaxRows * 32; + +struct Candidate +{ + float value; + int index; +}; + +// Total order of at::native GreaterOrNan (the order torch.max and torch.argmax use): a NaN beats every number and the +// lowest-index NaN wins; otherwise the larger value, and the lower index among equal values. +__device__ __forceinline__ bool isBetter(float v1, int i1, float v2, int i2) +{ + bool const nan1 = v1 != v1; + bool const nan2 = v2 != v2; + if (nan1 || nan2) + { + return nan1 && (!nan2 || i1 < i2); + } + return v1 > v2 || (v1 == v2 && i1 < i2); +} + +__device__ __forceinline__ void offer(Candidate& best, float value, int index) +{ + if (isBetter(value, index, best.value, best.index)) + { + best.value = value; + best.index = index; + } +} + +__device__ __forceinline__ Candidate warpArgmax(Candidate c) +{ +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + { + float const value = __shfl_xor_sync(0xFFFFFFFFU, c.value, offset); + int const index = __shfl_xor_sync(0xFFFFFFFFU, c.index, offset); + offer(c, value, index); + } + return c; +} + +// BF16 element 2k (low half) and 2k + 1 (high half) of a 32-bit word, widened exactly. +__device__ __forceinline__ float bf16Low(uint32_t word) +{ + return __uint_as_float(word << 16); +} + +__device__ __forceinline__ float bf16High(uint32_t word) +{ + return __uint_as_float(word & 0xFFFF0000U); +} + +__device__ __forceinline__ void offerVector(Candidate& best, uint4 const& v, int col) +{ + offer(best, bf16Low(v.x), col); + offer(best, bf16High(v.x), col + 1); + offer(best, bf16Low(v.y), col + 2); + offer(best, bf16High(v.y), col + 3); + offer(best, bf16Low(v.z), col + 4); + offer(best, bf16High(v.z), col + 5); + offer(best, bf16Low(v.w), col + 6); + offer(best, bf16High(v.w), col + 7); +} + +// This thread's best element of a logits row (vocabShard / 8 vectors of 8 BF16). +__device__ Candidate scanRow(uint4 const* row, int numVecs) +{ + Candidate best{-INFINITY, INT_MAX}; + for (int base = static_cast(threadIdx.x); base < numVecs; base += kScanBatch * kBlockSize) + { + uint4 vecs[kScanBatch]; +#pragma unroll + for (int k = 0; k < kScanBatch; ++k) + { + int const v = base + k * kBlockSize; + if (v < numVecs) + { + vecs[k] = row[v]; + } + } +#pragma unroll + for (int k = 0; k < kScanBatch; ++k) + { + int const v = base + k * kBlockSize; + if (v < numVecs) + { + offerVector(best, vecs[k], v * kNormVecElts); + } + } + } + return best; +} + +__device__ __forceinline__ bool isEmptyPacket(VolatilePackedLoad const& packet) +{ + return packet.words[0] == kEmptyWord || packet.words[1] == kEmptyWord || packet.words[2] == kEmptyWord + || packet.words[3] == kEmptyWord; +} + +__device__ __forceinline__ void namedBarrier(int id, int numThreads) +{ + asm volatile("bar.sync %0, %1;" ::"r"(id), "r"(numThreads) : "memory"); +} + +__device__ __forceinline__ float rsqrtApprox(float x) +{ + float y; + asm("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ uint32_t packBf16(float lo, float hi) +{ + __nv_bfloat162 const v = __floats2bfloat162_rn(lo, hi); + return *reinterpret_cast(&v); +} + +// (x * rstd) * (w + 0) of 8 elements, rounded to BF16. +__device__ __forceinline__ uint4 normalizeVector(uint4 const& x, uint4 const& w, float rstd) +{ + auto const scaled = [rstd](float xv, float wv) { return __fmul_rn(__fmul_rn(xv, rstd), __fadd_rn(wv, 0.F)); }; + uint4 y; + y.x = packBf16(scaled(bf16Low(x.x), bf16Low(w.x)), scaled(bf16High(x.x), bf16High(w.x))); + y.y = packBf16(scaled(bf16Low(x.y), bf16Low(w.y)), scaled(bf16High(x.y), bf16High(w.y))); + y.z = packBf16(scaled(bf16Low(x.z), bf16Low(w.z)), scaled(bf16High(x.z), bf16High(w.z))); + y.w = packBf16(scaled(bf16Low(x.w), bf16Low(w.w)), scaled(bf16High(x.w), bf16High(w.w))); + return y; +} + +__device__ __forceinline__ float sumOfSquares(uint4 const& v, float acc) +{ + uint32_t const words[4] = {v.x, v.y, v.z, v.w}; +#pragma unroll + for (int k = 0; k < 4; ++k) + { + float const lo = bf16Low(words[k]); + float const hi = bf16High(words[k]); + acc = __fadd_rn(acc, __fmul_rn(lo, lo)); + acc = __fadd_rn(acc, __fmul_rn(hi, hi)); + } + return acc; +} + +// RMSNorm of one row (src, or a zero row when null) by one 64-thread group, with the arithmetic of FlashInfer's +// RMSNormKernel: each thread sums the squares of its 96 elements in column order, the warp adds the partial sums in a +// butterfly (offsets 1, 2, ..., 16), the row sum is the first warp's total plus the second's; then +// rsqrt.approx(sum / H + eps) and (x * rstd) * (w + 0) in FP32, rounded to BF16. Columns [lo, hi) of the row (half +// coordinates) are stored at out[col - lo]. +__device__ void normRowWindow(uint4 const* src, __nv_bfloat16 const* weights, float eps, int lo, int hi, + __nv_bfloat16* out, int group, float (*groupSums)[2]) +{ + int const t = static_cast(threadIdx.x) % kNormThreads; + int const warpInGroup = t / 32; + int const lane = t % 32; + uint4 x[kNormVecs]; +#pragma unroll + for (int b = 0; b < kNormVecs; ++b) + { + x[b] = src != nullptr ? src[t + kNormThreads * b] : make_uint4(0U, 0U, 0U, 0U); + } + float acc = 0.F; +#pragma unroll + for (int b = 0; b < kNormVecs; ++b) + { + acc = sumOfSquares(x[b], acc); + } +#pragma unroll + for (int offset = 1; offset < 32; offset <<= 1) + { + acc = __fadd_rn(acc, __shfl_xor_sync(0xFFFFFFFFU, acc, offset)); + } + if (lane == 0) + { + groupSums[group][warpInGroup] = acc; + } + namedBarrier(1 + group, kNormThreads); + float const sum = __fadd_rn(groupSums[group][0], groupSums[group][1]); + float const rstd = rsqrtApprox(__fadd_rn(__fdiv_rn(sum, static_cast(kDraftSampleGlueHidden)), eps)); +#pragma unroll + for (int b = 0; b < kNormVecs; ++b) + { + int const col = kNormVecElts * t + kNormThreads * kNormVecElts * b; + if (col >= lo && col < hi) + { + uint4 const w = *reinterpret_cast(weights + col); + *reinterpret_cast(out + (col - lo)) = normalizeVector(x[b], w, rstd); + } + } + // The group's sums are rewritten by its next row. + namedBarrier(1 + group, kNormThreads); +} + +// Inclusive warp prefix sum. +__device__ __forceinline__ int64_t warpInclusiveSum(int64_t value) +{ + int const lane = static_cast(threadIdx.x) % 32; +#pragma unroll + for (int offset = 1; offset < 32; offset <<= 1) + { + int64_t const other = __shfl_up_sync(0xFFFFFFFFU, value, offset); + if (lane >= offset) + { + value += other; + } + } + return value; +} + +// Position ids, sequence and KV lengths of the next step, then the DSA decode metadata of fused_dsa_decode_metadata +// (lane r handles request r). One warp. +__device__ void advanceStep(DraftSampleGlueParams const& p) +{ + int const lane = static_cast(threadIdx.x) % 32; + bool const active = lane < p.numRows; + int32_t seq = 0; + int32_t kv = 0; + if (active) + { + if (p.positionsOut != nullptr) + { + int64_t const src = p.gatherIds != nullptr ? p.gatherIds[lane] : lane; + if (p.positions64) + { + static_cast(p.positionsOut)[lane] = static_cast(p.positionsIn)[src] + 1; + } + else + { + static_cast(p.positionsOut)[lane] = static_cast(p.positionsIn)[src] + 1; + } + } + if (p.seqLens != nullptr) + { + seq = p.fillSeqLens ? 1 : p.seqLens[lane]; + if (p.fillSeqLens) + { + p.seqLens[lane] = 1; + } + } + if (p.kvLens != nullptr) + { + kv = p.kvLens[lane]; + if (p.kvUpdate == DraftSampleGlueKvUpdate::kIncrement) + { + kv += 1; + p.kvLens[lane] = kv; + } + else if (p.kvUpdate == DraftSampleGlueKvUpdate::kRewind) + { + kv -= p.draftLen - p.kvRewindAccepted[lane]; + p.kvLens[lane] = kv; + } + } + } + if (p.reqIdxPerToken == nullptr) + { + return; + } + int64_t const seq64 = active ? seq : 0; + int64_t const kv64 = active ? kv : 0; + int64_t const seqIncl = warpInclusiveSum(seq64); + int64_t const kvIncl = warpInclusiveSum(kv64); + int64_t const cachedIncl = warpInclusiveSum(kv64 - seq64); + if (lane == 0) + { + p.genKvIndptr[0] = 0; + p.genCachedTokenIndptr[0] = 0; + } + if (!active) + { + return; + } + p.genKvIndptr[1 + lane] = kvIncl; + p.genCachedTokenIndptr[1 + lane] = cachedIncl; + int64_t const flatStart = seqIncl - seq64; + int64_t const startPos = kv64 - seq64; + int64_t const tokens = seq64 < p.maxQueryLanes ? seq64 : p.maxQueryLanes; + for (int64_t j = 0; j < tokens; ++j) + { + int64_t const g = flatStart + j; + if (g >= p.numTokens) + { + break; + } + int64_t const globalPos = startPos + j; + // Truncating division and a non-negative remainder, then the clamp, as the Triton kernel computes them. + int64_t blk = globalPos / p.tokensPerBlock; + int64_t pos = globalPos % p.tokensPerBlock; + pos = pos < 0 ? pos + p.tokensPerBlock : pos; + blk = blk < 0 ? 0 : (blk > p.maxBlocks - 1 ? p.maxBlocks - 1 : blk); + int64_t const blockId = p.blockOffsets[lane * p.blockOffsetsStride0 + blk * p.blockOffsetsStride1]; + p.reqIdxPerToken[g] = lane; + p.slotMappingFp8[g] = blockId * p.blockStride + pos * p.dataBytesPerToken; + p.slotMappingScale[g] = blockId * p.blockStride + p.scaleBaseOffset + pos * p.scaleSize; + } +} + +__global__ void __launch_bounds__(kBlockSize) draftSampleGlueKernel(DraftSampleGlueParams const p) +{ + __shared__ __align__(16) __nv_bfloat16 sWeights[2][kDraftSampleGlueHidden]; + __shared__ Candidate sCandidates[kDraftSampleGlueMaxRows][kNumWarps]; + __shared__ Candidate sSlots[kMaxSlots]; + __shared__ int sTokens[kDraftSampleGlueMaxRows]; + __shared__ float sGroupSums[kNumNormGroups][2]; + __shared__ uint4 sFlags[2]; + + int const tid = static_cast(threadIdx.x); + int const warp = tid / 32; + int const lane = tid % 32; + bool const glue = p.glueOut != nullptr; + // The halves of the concatenated row that the eh_proj input window overlaps. + bool const embedHalf = glue && p.colLo < kDraftSampleGlueHidden; + bool const hiddenHalf = glue && p.colHi > kDraftSampleGlueHidden; + + // The norm weights are constants: copy the window columns before the grid dependency wait. + if (glue) + { + for (int h = 0; h < 2; ++h) + { + if ((h == 0 && !embedHalf) || (h == 1 && !hiddenHalf)) + { + continue; + } + int const lo = max(p.colLo - h * kDraftSampleGlueHidden, 0); + int const hi = min(p.colHi - h * kDraftSampleGlueHidden, kDraftSampleGlueHidden); + auto const* src = reinterpret_cast(h == 0 ? p.enormWeight : p.hnormWeight); + auto* dst = reinterpret_cast(sWeights[h]); + for (int v = lo / kNormVecElts + tid; v < hi / kNormVecElts; v += kBlockSize) + { + dst[v] = src[v]; + } + } + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + // The logits come from the preceding kernel (the draft LM head). Dependents read the outputs only after their + // own wait. + cudaGridDependencySynchronize(); + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + if (tid == 0) + { + sFlags[0] = reinterpret_cast(p.exchangeFlags)[0]; + sFlags[1] = reinterpret_cast(p.exchangeFlags)[1]; + } + + // Local argmax of every row: all threads scan, each warp reduces, warp 0 combines below. + int const numVecs = p.vocabShard / kNormVecElts; + for (int row = 0; row < p.numRows; ++row) + { + Candidate const c + = warpArgmax(scanRow(reinterpret_cast(p.logits + row * p.logitsStride), numVecs)); + if (lane == 0) + { + sCandidates[row][warp] = c; + } + } + __syncthreads(); + + uint4 const flags = sFlags[0]; + uint32_t const curIndex = flags.x; + uint32_t const dirtyIndex = flags.y; + uint32_t const bytesPerBuffer = flags.z; + uint32_t const dirtyStages = flags.w; + int const numSlots = p.numRows * p.tpSize; + auto* exchangeLocal = static_cast(p.exchangeLocal); + + if (warp == 0) + { + // Publish this rank's pair of every row through the multicast address: slot [row][rank] of the current buffer. + auto* curMcast = reinterpret_cast(static_cast(p.exchangeMcast) + curIndex * bytesPerBuffer); + for (int row = 0; row < p.numRows; ++row) + { + Candidate c = lane < kNumWarps ? sCandidates[row][lane] : Candidate{-INFINITY, INT_MAX}; + c = warpArgmax(c); + if (lane == 0) + { + int const globalIndex = c.index + p.rank * p.vocabShard; + curMcast[row * p.tpSize + p.rank] + = make_uint4(static_cast(globalIndex), __float_as_uint(c.value), 0U, 0U); + } + } + // Poll the local copy until every rank's pair has landed. + auto const* curLocal = exchangeLocal + curIndex * bytesPerBuffer; + for (int s = lane; s < numSlots; s += 32) + { + VolatilePackedLoad packet; + do + { + packet = loadPackedVolatile(curLocal + s * kDraftSampleGluePacketBytes); + } while (isEmptyPacket(packet)); + sSlots[s] = Candidate{__uint_as_float(packet.words[1]), static_cast(packet.words[0])}; + } + __syncwarp(); + // Global argmax per row: the ranks' pairs in rank order (the lower rank holds the lower indices). + if (lane < p.numRows) + { + Candidate best = sSlots[lane * p.tpSize]; + for (int r = 1; r < p.tpSize; ++r) + { + offer(best, sSlots[lane * p.tpSize + r].value, sSlots[lane * p.tpSize + r].index); + } + sTokens[lane] = best.index; + p.draftTokens[lane * p.numDraftTokens + p.draftStep] = best.index; + } + } + else if (warp == 1) + { + // Empty the buffer of the previous call; the peers write it again two calls later. + if (dirtyStages > 0) + { + auto* dirty = reinterpret_cast(exchangeLocal + dirtyIndex * bytesPerBuffer); + int const dirtyPackets = static_cast(sFlags[1].x / kDraftSampleGluePacketBytes); + for (int i = lane; i < dirtyPackets; i += 32) + { + dirty[i] = make_uint4(kEmptyWord, kEmptyWord, kEmptyWord, kEmptyWord); + } + } + advanceStep(p); + } + else if (hiddenHalf) + { + // hnorm of the hidden rows does not need the tokens: groups 1..7 (warps 2..15) run it during the poll. + int const group = warp / 2; + int const lo = max(p.colLo - kDraftSampleGlueHidden, 0); + int const hi = min(p.colHi - kDraftSampleGlueHidden, kDraftSampleGlueHidden); + int const outShift = kDraftSampleGlueHidden + lo - p.colLo; + for (int row = group - 1; row < p.numRows; row += kNumNormGroups - 1) + { + normRowWindow(reinterpret_cast(p.hidden + row * p.hiddenStride), sWeights[1], p.eps, lo, hi, + p.glueOut + row * p.glueOutStride + outShift, group, sGroupSums); + } + } + __syncthreads(); + + if (embedHalf) + { + // enorm of the token embeddings: an id outside [0, vocabSize) gives a zero row, as the Embedding module does. + int const group = warp / 2; + int const lo = p.colLo; + int const hi = min(p.colHi, kDraftSampleGlueHidden); + for (int row = group; row < p.numRows; row += kNumNormGroups) + { + int64_t const token = sTokens[row]; + bool const valid = token >= 0 && token < p.vocabSize; + normRowWindow( + valid ? reinterpret_cast(p.embedTable + token * kDraftSampleGlueHidden) : nullptr, + sWeights[0], p.eps, lo, hi, p.glueOut + row * p.glueOutStride, group, sGroupSums); + } + } + + if (warp == 1 && p.nextNewTokens != nullptr && lane < p.numRows) + { + // Output rows after the last draft step: the last accepted token, then the draft tokens of every step. + int32_t* out = p.nextNewTokens + lane * (p.numDraftTokens + 1); + out[0] = p.acceptedTokens[p.batchIndices[lane] * p.acceptedTokensStride + p.numAccepted[lane] - 1]; + for (int j = 0; j < p.numDraftTokens; ++j) + { + out[1 + j] = j == p.draftStep ? sTokens[lane] : p.draftTokens[lane * p.numDraftTokens + j]; + } + } + + if (tid == 0) + { + // Rotate the Lamport flags as LamportFlags::update does: the current buffer becomes the dirty one. + auto* flagPtr = reinterpret_cast(p.exchangeFlags); + flagPtr[0] = make_uint4((curIndex + 1) % 3, curIndex, bytesPerBuffer, 1U); + flagPtr[1] = make_uint4(static_cast(numSlots * kDraftSampleGluePacketBytes), 0U, 0U, 0U); + p.exchangeFlags[8] = 0U; + } +} + +// The draft step with the next step's eh_proj, all-reduce and input_layernorm (draftEhprojTableKernel). The cluster +// covers the hidden row: thread t < kTableSliceVecs of CTA c owns columns [c * kTableSliceElts + 8 t, + 8). +constexpr int kTableCtas = kDraftEhprojTableClusterSize; +constexpr int kTableSliceElts = kDraftSampleGlueHidden / kTableCtas; +constexpr int kTableSliceVecs = kTableSliceElts / kNormVecElts; +constexpr int kTableSliceBytes = kTableSliceElts * static_cast(sizeof(__nv_bfloat16)); +constexpr int kTableVecWarps = kTableSliceVecs / 32; +// After the vector warps: a warp that polls the (index, value) packets and scans, then a warp that updates the next +// step's state (CTA 0) or empties the previous call's buffer (the other CTAs) beside the exchange. +constexpr int kTableExchangeWarp = kTableVecWarps; +constexpr int kTableStateWarp = kTableVecWarps + 1; +constexpr int kTableMainThreads = kTableStateWarp * 32; +constexpr int kTableThreads = kTableMainThreads + 32; +constexpr int kTableMainWarps = kTableMainThreads / 32; +// Rank slots of the kernel (its template argument): eight up to TP8, kDraftEhprojTableMaxRanks above. +constexpr int kTableTp8RankSlots = kTableCtas; +// Logit vectors each scanning thread loads before it compares them. +constexpr int kTableScanBatch = 4; +// Slot layout: the packet, the BF16 token-table row, then the rank's FP32 hidden-half outputs. +constexpr int kTableRowOffset = kDraftSampleGluePacketBytes; +constexpr int kTableHiddenOffset = kTableRowOffset + 2 * kDraftSampleGlueHidden; +// Named barrier of the main warps (the state warp runs on its own). +constexpr int kTableMainBarrier = 1; +static_assert(kTableSliceVecs % 32 == 0, "the slice must be whole warps of vectors"); +static_assert(kTableSliceElts * kTableCtas == kDraftSampleGlueHidden, "the cluster must cover the row"); + +__device__ __forceinline__ uint4 sanitizeWords(uint4 v) +{ + // A word with every bit set is the empty word; such a payload word (two NaN halves) is sent as another NaN. + v.x = v.x == kEmptyWord ? kEmptyWord - 1U : v.x; + v.y = v.y == kEmptyWord ? kEmptyWord - 1U : v.y; + v.z = v.z == kEmptyWord ? kEmptyWord - 1U : v.z; + v.w = v.w == kEmptyWord ? kEmptyWord - 1U : v.w; + return v; +} + +__device__ __forceinline__ float warpSum(float v) +{ +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + { + v = __fadd_rn(v, __shfl_xor_sync(0xFFFFFFFFU, v, offset)); + } + return v; +} + +__device__ __forceinline__ void clusterArriveRelaxed() +{ + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); +} + +__device__ __forceinline__ void clusterArrive() +{ + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); +} + +__device__ __forceinline__ void clusterWait() +{ + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); +} + +template +__global__ void __launch_bounds__(kTableThreads) draftEhprojTableKernel(DraftEhprojTableParams const p) +{ + // sCtaBest is read by the other CTAs of the cluster after the first cluster barrier; sClusterSums is written by + // every CTA (slot = the writer's rank in the cluster) before the second one and read locally after it. The rest + // is shared by the main warps through the named barrier. + __shared__ Candidate sWarpBest[kTableMainWarps]; + __shared__ float sWarpSums[kTableMainWarps]; + __shared__ Candidate sCtaBest[kDraftSampleGlueMaxRows]; + __shared__ Candidate sRankBest[kDraftSampleGlueMaxRows]; + __shared__ float sClusterSums[kDraftSampleGlueMaxRows][kTableCtas]; + __shared__ float sRstd[kDraftSampleGlueMaxRows]; + __shared__ Candidate sPairs[kDraftSampleGlueMaxRows][kMaxRanks]; + __shared__ __align__(16) __nv_bfloat16 sX[kDraftSampleGlueMaxRows][kTableSliceElts]; + __shared__ uint4 sFlags; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + DraftSampleGlueParams const& s = p.step; + int const tid = static_cast(threadIdx.x); + int const warp = tid / 32; + int const lane = tid % 32; + int const cta = static_cast(cluster.block_rank()); + int const tp = s.tpSize; + int const colBase = cta * kTableSliceElts; + bool const vecThread = tid < kTableSliceVecs; + // The rank whose share of the hidden-half outputs (H / tp columns) holds this vector thread's eight columns, and + // their offset in that share: a CTA's columns lie in one share up to TP8 and span two shares at TP16. + int const shareCols = kDraftSampleGlueHidden / tp; + int const vecCol = colBase + kNormVecElts * tid; + int const hiddenOwner = vecCol / shareCols; + int const shareOffset = vecCol - hiddenOwner * shareCols; + int const slotBytes = draftEhprojTableSlotBytes(tp); + + // The norm weight is a constant: load it before the grid dependency wait. + uint4 normWeight = make_uint4(0U, 0U, 0U, 0U); + if (vecThread) + { + normWeight = reinterpret_cast(p.normWeight + colBase)[tid]; + } + + // The logits come from the preceding kernel (the draft LM head), the hidden-half outputs from a kernel this launch + // depends on through an event. Dependents read the outputs only after their own wait. + cudaGridDependencySynchronize(); + cudaTriggerProgrammaticLaunchCompletion(); + + if (warp == kTableStateWarp) + { + // Arrive for the first cluster barrier before the state work, which nothing in the cluster waits for. + clusterArriveRelaxed(); + if (cta == 0) + { + advanceStep(s); + } + else + { + // Empty this CTA's share of the previous call's buffer; the peers write it again two calls later. + uint4 const flags = reinterpret_cast(s.exchangeFlags)[0]; + uint4 const sizes = reinterpret_cast(s.exchangeFlags)[1]; + if (flags.w > 0) + { + auto* dirty = reinterpret_cast( + static_cast(s.exchangeLocal) + static_cast(flags.y) * flags.z); + int const vectors = static_cast(sizes.x / sizeof(uint4)); + for (int i = (cta - 1) * 32 + lane; i < vectors; i += (kTableCtas - 1) * 32) + { + dirty[i] = make_uint4(kEmptyWord, kEmptyWord, kEmptyWord, kEmptyWord); + } + } + } + clusterWait(); + // The second (last) cluster barrier. + clusterArrive(); + clusterWait(); + return; + } + + if (tid == 0) + { + sFlags = reinterpret_cast(s.exchangeFlags)[0]; + } + + // Per row: this CTA's best logit (the CTAs scan consecutive chunks of the shard). + int const numVecs = s.vocabShard / kNormVecElts; + int const chunk = (numVecs + kTableCtas - 1) / kTableCtas; + int const vecBegin = min(cta * chunk, numVecs); + int const vecEnd = min(vecBegin + chunk, numVecs); + for (int row = 0; row < s.numRows; ++row) + { + auto const* logitsRow = reinterpret_cast(s.logits + row * s.logitsStride); + Candidate best{-INFINITY, INT_MAX}; + for (int base = vecBegin + tid; base < vecEnd; base += kTableScanBatch * kTableMainThreads) + { + uint4 vecs[kTableScanBatch]; +#pragma unroll + for (int k = 0; k < kTableScanBatch; ++k) + { + int const v = base + k * kTableMainThreads; + if (v < vecEnd) + { + vecs[k] = logitsRow[v]; + } + } +#pragma unroll + for (int k = 0; k < kTableScanBatch; ++k) + { + int const v = base + k * kTableMainThreads; + if (v < vecEnd) + { + offerVector(best, vecs[k], v * kNormVecElts); + } + } + } + best = warpArgmax(best); + if (lane == 0) + { + sWarpBest[warp] = best; + } + namedBarrier(kTableMainBarrier, kTableMainThreads); + Candidate ctaBest = sWarpBest[0]; + for (int w = 1; w < kTableMainWarps; ++w) + { + offer(ctaBest, sWarpBest[w].value, sWarpBest[w].index); + } + if (tid == 0) + { + sCtaBest[row] = ctaBest; + } + // The winner's token-table row is one of the cluster's candidates: pull this CTA's candidate row into L2 while + // the cluster combines them. + if (vecThread && ctaBest.index >= 0 && ctaBest.index < s.vocabShard) + { + auto const* row = p.tokenTable + static_cast(ctaBest.index) * kDraftSampleGlueHidden; + char const* line = reinterpret_cast(row) + tid * 128; + asm volatile("prefetch.global.L2 [%0];" ::"l"(line) : "memory"); + } + // sWarpBest is rewritten for the next row only after this barrier. + namedBarrier(kTableMainBarrier, kTableMainThreads); + } + + uint4 const flags = sFlags; + uint32_t const curIndex = flags.x; + uint32_t const bytesPerBuffer = flags.z; + char* const curMcast = static_cast(s.exchangeMcast) + static_cast(curIndex) * bytesPerBuffer; + char const* const curLocal + = static_cast(s.exchangeLocal) + static_cast(curIndex) * bytesPerBuffer; + + // This rank's hidden-half outputs do not depend on the exchange: publish them first. + if (vecThread && hiddenOwner == s.rank) + { + for (int row = 0; row < s.numRows; ++row) + { + auto const* h = reinterpret_cast(p.hiddenHalf + row * p.hiddenHalfStride + shareOffset); + char* const slot = curMcast + static_cast(row * tp + s.rank) * slotBytes; + auto* dst = reinterpret_cast(slot + kTableHiddenOffset + shareOffset * sizeof(float)); + uint4 const h0 = h[0]; + uint4 const h1 = h[1]; + dst[0] = sanitizeWords(h0); + dst[1] = sanitizeWords(h1); + } + } + + clusterArrive(); + clusterWait(); + + // This rank's candidate of every row: the CTAs' candidates in CTA order. These are the last reads of another CTA's + // shared memory; every CTA makes them before it arrives at the second cluster barrier. + if (tid < s.numRows) + { + Candidate rankBest = *cluster.map_shared_rank(&sCtaBest[tid], 0); + for (int c = 1; c < kTableCtas; ++c) + { + Candidate const other = *cluster.map_shared_rank(&sCtaBest[tid], c); + offer(rankBest, other.value, other.index); + } + sRankBest[tid] = rankBest; + } + namedBarrier(kTableMainBarrier, kTableMainThreads); + + for (int row = 0; row < s.numRows; ++row) + { + Candidate const rankBest = sRankBest[row]; + char* const slot = curMcast + static_cast(row * tp + s.rank) * slotBytes; + if (vecThread) + { + // This CTA's columns of the token-table row of the rank's candidate. + uint4 const q = reinterpret_cast( + p.tokenTable + static_cast(rankBest.index) * kDraftSampleGlueHidden + colBase)[tid]; + reinterpret_cast(slot + kTableRowOffset + cta * kTableSliceBytes)[tid] = sanitizeWords(q); + } + else if (cta == 0 && tid == kTableExchangeWarp * 32) + { + int const globalIndex = rankBest.index + s.rank * s.vocabShard; + *reinterpret_cast(slot) + = make_uint4(static_cast(globalIndex), __float_as_uint(rankBest.value), 0U, 0U); + } + + // Poll the local copy: the ranks' rows of this CTA's columns and the two vectors of hidden-half outputs that + // cover them (vector threads), the ranks' packets (exchange warp). Every round has all its loads in flight. + char const* const rowLocal = curLocal + static_cast(row * tp) * slotBytes; + VolatilePackedLoad rows[kMaxRanks]; + VolatilePackedLoad hidden0; + VolatilePackedLoad hidden1; + if (vecThread) + { + char const* const rowSlice = rowLocal + kTableRowOffset + cta * kTableSliceBytes + tid * sizeof(uint4); + char const* const hiddenPtr = rowLocal + static_cast(hiddenOwner) * slotBytes + kTableHiddenOffset + + shareOffset * sizeof(float); + bool empty = true; + while (empty) + { +#pragma unroll + for (int r = 0; r < kMaxRanks; ++r) + { + if (r < tp) + { + rows[r] = loadPackedVolatile(rowSlice + static_cast(r) * slotBytes); + } + } + hidden0 = loadPackedVolatile(hiddenPtr); + hidden1 = loadPackedVolatile(hiddenPtr + sizeof(uint4)); + empty = isEmptyPacket(hidden0) || isEmptyPacket(hidden1); +#pragma unroll + for (int r = 0; r < kMaxRanks; ++r) + { + if (r < tp) + { + empty = empty || isEmptyPacket(rows[r]); + } + } + } + } + else if (warp == kTableExchangeWarp && lane < tp) + { + VolatilePackedLoad packet; + do + { + packet = loadPackedVolatile(rowLocal + static_cast(lane) * slotBytes); + } while (isEmptyPacket(packet)); + sPairs[row][lane] = Candidate{__uint_as_float(packet.words[1]), static_cast(packet.words[0])}; + } + namedBarrier(kTableMainBarrier, kTableMainThreads); + + // The global argmax: the ranks' packets in rank order (the lower rank holds the lower indices). + int winner = 0; + Candidate best = sPairs[row][0]; + for (int r = 1; r < tp; ++r) + { + if (isBetter(sPairs[row][r].value, sPairs[row][r].index, best.value, best.index)) + { + best = sPairs[row][r]; + winner = r; + } + } + if (cta == 0 && tid == 0) + { + s.draftTokens[row * s.numDraftTokens + s.draftStep] = best.index; + } + + float sumSq = 0.F; + if (vecThread) + { + uint32_t qw[4] = {rows[0].words[0], rows[0].words[1], rows[0].words[2], rows[0].words[3]}; +#pragma unroll + for (int r = 1; r < kMaxRanks; ++r) + { +#pragma unroll + for (int k = 0; k < 4; ++k) + { + qw[k] = r == winner ? rows[r].words[k] : qw[k]; + } + } + // x = token-table row + hidden-half output: one FP32 add and one BF16 rounding per element. + uint32_t const hw[8] = {hidden0.words[0], hidden0.words[1], hidden0.words[2], hidden0.words[3], + hidden1.words[0], hidden1.words[1], hidden1.words[2], hidden1.words[3]}; + uint32_t xw[4]; +#pragma unroll + for (int k = 0; k < 4; ++k) + { + xw[k] = packBf16(__fadd_rn(bf16Low(qw[k]), __uint_as_float(hw[2 * k])), + __fadd_rn(bf16High(qw[k]), __uint_as_float(hw[2 * k + 1]))); + } + uint4 const x = make_uint4(xw[0], xw[1], xw[2], xw[3]); + reinterpret_cast(sX[row])[tid] = x; + reinterpret_cast(p.residualOut + row * p.outStride + colBase)[tid] = x; + sumSq = sumOfSquares(x, 0.F); + } + sumSq = warpSum(sumSq); + if (lane == 0) + { + sWarpSums[warp] = sumSq; + } + namedBarrier(kTableMainBarrier, kTableMainThreads); + if (tid < kTableCtas) + { + float sum = sWarpSums[0]; + for (int w = 1; w < kTableVecWarps; ++w) + { + sum = __fadd_rn(sum, sWarpSums[w]); + } + *cluster.map_shared_rank(&sClusterSums[row][cta], tid) = sum; + } + // sPairs and sWarpSums are rewritten for the next row only after this barrier. + namedBarrier(kTableMainBarrier, kTableMainThreads); + } + + // Every CTA has sent its sums of squares and read the others' candidates: after this barrier no CTA touches + // another's shared memory. + clusterArrive(); + clusterWait(); + + // input_layernorm: the row's sum of squares over the cluster in CTA order. + if (tid < s.numRows) + { + float sum = sClusterSums[tid][0]; + for (int c = 1; c < kTableCtas; ++c) + { + sum = __fadd_rn(sum, sClusterSums[tid][c]); + } + sRstd[tid] = rsqrtApprox(__fadd_rn(__fdiv_rn(sum, static_cast(kDraftSampleGlueHidden)), p.normEps)); + } + namedBarrier(kTableMainBarrier, kTableMainThreads); + if (vecThread) + { + for (int row = 0; row < s.numRows; ++row) + { + uint4 const x = reinterpret_cast(sX[row])[tid]; + reinterpret_cast(p.normedOut + row * p.outStride + colBase)[tid] + = normalizeVector(x, normWeight, sRstd[row]); + } + } + + if (cta == 0 && tid == 0) + { + // Every CTA read the flags before the first cluster barrier, every state warp before the second: rotate them + // as LamportFlags::update does. + auto* flagPtr = reinterpret_cast(s.exchangeFlags); + flagPtr[0] = make_uint4((curIndex + 1) % 3, curIndex, bytesPerBuffer, 1U); + flagPtr[1] = make_uint4(static_cast(s.numRows * tp * slotBytes), 0U, 0U, 0U); + s.exchangeFlags[8] = 0U; + } +#endif +} + +} // namespace + +void invokeDraftSampleGlue(DraftSampleGlueParams const& params, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(params.numRows > 0 && params.numRows <= kDraftSampleGlueMaxRows, + "draft sample glue serves 1 to %d rows, got %d", kDraftSampleGlueMaxRows, params.numRows); + TLLM_CHECK_WITH_INFO(params.tpSize >= 1 && params.numRows * params.tpSize <= kDraftSampleGlueMaxRows * 32, + "draft sample glue: %d rows x %d ranks exceed the exchange slots", params.numRows, params.tpSize); + TLLM_CHECK_WITH_INFO(params.vocabShard > 0 && params.vocabShard % 8 == 0, + "draft sample glue needs a vocabulary shard of a multiple of 8 columns, got %d", params.vocabShard); + TLLM_CHECK_WITH_INFO(params.glueOut == nullptr + || (params.colLo >= 0 && params.colLo < params.colHi && params.colHi <= 2 * kDraftSampleGlueHidden + && params.colLo % 8 == 0 && params.colHi % 8 == 0), + "draft sample glue: invalid eh_proj input window [%d, %d)", params.colLo, params.colHi); + TLLM_CHECK_WITH_INFO( + params.glueOut == nullptr || params.colLo >= kDraftSampleGlueHidden || params.embedTable != nullptr, + "draft sample glue: the window needs the embedding table"); + TLLM_CHECK_WITH_INFO( + params.glueOut == nullptr || params.colHi <= kDraftSampleGlueHidden || params.hidden != nullptr, + "draft sample glue: the window needs the hidden rows"); + TLLM_CHECK_WITH_INFO(params.reqIdxPerToken == nullptr || (params.seqLens != nullptr && params.kvLens != nullptr), + "draft sample glue: the DSA metadata needs the sequence and KV lengths"); + tensorrt_llm::common::launchWithPdlWhenEnabled( + "draftSampleGlue", draftSampleGlueKernel, dim3(1), dim3(kBlockSize), 0, stream, params); +} + +void invokeDraftEhprojTable(DraftEhprojTableParams const& params, cudaStream_t stream) +{ + DraftSampleGlueParams const& s = params.step; + TLLM_CHECK_WITH_INFO(s.numRows > 0 && s.numRows <= kDraftSampleGlueMaxRows, + "draft eh_proj table serves 1 to %d rows, got %d", kDraftSampleGlueMaxRows, s.numRows); + TLLM_CHECK_WITH_INFO(draftEhprojTableSupportsTpSize(s.tpSize), + "draft eh_proj table serves 2 to %d ranks with whole 8-column shares of the hidden size, got %d", + kDraftEhprojTableMaxRanks, s.tpSize); + TLLM_CHECK_WITH_INFO(s.vocabShard > 0 && s.vocabShard % 8 == 0 && s.logitsStride % 8 == 0, + "draft eh_proj table needs 16-byte aligned rows of a multiple of 8 logits, got %d", s.vocabShard); + TLLM_CHECK_WITH_INFO(params.hiddenHalf != nullptr && params.hiddenHalfStride % 8 == 0 + && reinterpret_cast(params.hiddenHalf) % 32 == 0, + "draft eh_proj table needs 32-byte aligned rows of FP32 hidden-half outputs"); + TLLM_CHECK_WITH_INFO(s.glueOut == nullptr && s.nextNewTokens == nullptr, + "draft eh_proj table: no eh_proj input window and no next-new-tokens rows"); + TLLM_CHECK_WITH_INFO(params.tokenTable != nullptr && params.normWeight != nullptr && params.normedOut != nullptr + && params.residualOut != nullptr && params.outStride % 8 == 0, + "draft eh_proj table: missing token table, norm weight or 16-byte aligned output rows"); + TLLM_CHECK_WITH_INFO(s.reqIdxPerToken == nullptr || (s.seqLens != nullptr && s.kvLens != nullptr), + "draft eh_proj table: the DSA metadata needs the sequence and KV lengths"); + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kDraftEhprojTableClusterSize; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL() ? 1 : 0; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(kDraftEhprojTableClusterSize); + config.blockDim = dim3(kTableThreads); + config.dynamicSmemBytes = 0; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = 2; + // Up to TP8 the kernel keeps eight rank slots (registers of the polled rows); TP16 takes sixteen. + auto* const kernel = s.tpSize <= kTableTp8RankSlots ? &draftEhprojTableKernel + : &draftEhprojTableKernel; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, params)); +} + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.h new file mode 100644 index 000000000000..25920a5f26fb --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.h @@ -0,0 +1,188 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +//! Rows (requests) one launch serves. +constexpr int kDraftSampleGlueMaxRows = 8; +//! Hidden size of the eh_proj input glue (FlashInfer RMSNormKernel geometry: 64 threads x 8-element vectors x 12). +constexpr int kDraftSampleGlueHidden = 6144; +//! Bytes of one exchanged (index, value) packet. +constexpr int kDraftSampleGluePacketBytes = 16; + +//! How the next step's KV lengths follow from the current ones. +enum class DraftSampleGlueKvUpdate : int +{ + kNone = 0, //!< unchanged + kIncrement = 1, //!< kvLens[r] += 1 + kRewind = 2, //!< kvLens[r] -= draftLen - kvRewindAccepted[r] +}; + +//! One MTP draft step after the vocabulary-sharded draft LM head, in one CTA: +//! - the greedy draft token: the local argmax of this rank's logits shard, exchanged with the other ranks through a +//! dedicated MNNVL Lamport workspace (layout and flags of the MNNVL all-gather), and the global argmax of the pairs; +//! - the state of the next draft step: position ids, sequence and KV lengths, the DSA decode metadata of +//! fused_dsa_decode_metadata; +//! - the next step's eh_proj input: this rank's columns of cat([enorm(embed(token)), hnorm(hidden)], -1); +//! - after the last draft step, the next-new-tokens rows of the output. +//! Every part except the token is optional (skipped when its output pointer is null). +struct DraftSampleGlueParams +{ + // Draft logits of this rank, [numRows, vocabShard] BF16 with row stride logitsStride (elements). + __nv_bfloat16 const* logits{nullptr}; + int64_t logitsStride{0}; + int numRows{0}; + int vocabShard{0}; + int rank{0}; + int tpSize{1}; + + // Lamport workspace: three buffers of kDraftSampleGlueMaxRows * tpSize packets, the local unicast view, the + // multicast view and the nine 32-bit flags of the MNNVL all-gather. + void* exchangeLocal{nullptr}; + void* exchangeMcast{nullptr}; + uint32_t* exchangeFlags{nullptr}; + + // Draft tokens [numRows, numDraftTokens] int32; this step writes column draftStep. + int32_t* draftTokens{nullptr}; + int numDraftTokens{1}; + int draftStep{0}; + + // Next step's eh_proj input (skipped when glueOut is null): columns [colLo, colHi) of + // cat([enorm(embedTable[token]), hnorm(hidden)], -1), stored at glueOut[row * glueOutStride + col - colLo]. + __nv_bfloat16 const* embedTable{nullptr}; + int64_t vocabSize{0}; + __nv_bfloat16 const* hidden{nullptr}; + int64_t hiddenStride{0}; + __nv_bfloat16 const* enormWeight{nullptr}; + __nv_bfloat16 const* hnormWeight{nullptr}; + float eps{0.F}; + int colLo{0}; + int colHi{0}; + __nv_bfloat16* glueOut{nullptr}; + int64_t glueOutStride{0}; + + // Next step's position ids (skipped when positionsOut is null): + // positionsOut[r] = positionsIn[gatherIds ? gatherIds[r] : r] + 1, int32 or int64. + void const* positionsIn{nullptr}; + void* positionsOut{nullptr}; + bool positions64{false}; + int64_t const* gatherIds{nullptr}; + + // Sequence and KV lengths of the next step: seqLens[r] = 1 when fillSeqLens; kvLens[r] updated as kvUpdate says. + int32_t* seqLens{nullptr}; + bool fillSeqLens{false}; + int32_t* kvLens{nullptr}; + DraftSampleGlueKvUpdate kvUpdate{DraftSampleGlueKvUpdate::kNone}; + int32_t const* kvRewindAccepted{nullptr}; + int draftLen{0}; + + // DSA decode metadata of the next step (skipped when reqIdxPerToken is null), from the sequence and KV lengths + // above: what fused_dsa_decode_metadata writes for numRows requests and numTokens tokens. maxQueryLanes is its + // per-request token lanes (the power of two at or above the maximum query length). + int32_t const* blockOffsets{nullptr}; + int64_t blockOffsetsStride0{0}; + int64_t blockOffsetsStride1{1}; + int maxBlocks{0}; + int numTokens{0}; + int maxQueryLanes{1}; + int tokensPerBlock{0}; + int64_t dataBytesPerToken{0}; + int64_t scaleSize{0}; + int64_t blockStride{0}; + int64_t scaleBaseOffset{0}; + int32_t* reqIdxPerToken{nullptr}; + int64_t* slotMappingFp8{nullptr}; + int64_t* slotMappingScale{nullptr}; + int64_t* genKvIndptr{nullptr}; + int64_t* genCachedTokenIndptr{nullptr}; + + // Output rows after the last draft step (skipped when nextNewTokens is null), [numRows, numDraftTokens + 1] + // int32: nextNewTokens[r, 0] = acceptedTokens[batchIndices[r], numAccepted[r] - 1], then the draft tokens. + int32_t const* acceptedTokens{nullptr}; + int64_t acceptedTokensStride{0}; + int32_t const* batchIndices{nullptr}; + int32_t const* numAccepted{nullptr}; + int32_t* nextNewTokens{nullptr}; +}; + +void invokeDraftSampleGlue(DraftSampleGlueParams const& params, cudaStream_t stream); + +//! CTAs of the cluster that runs invokeDraftEhprojTable; each owns kDraftSampleGlueHidden / 8 hidden columns. +constexpr int kDraftEhprojTableClusterSize = 8; +//! Largest TP size of invokeDraftEhprojTable (2, 4, 8 or 16 ranks: each rank's share of the hidden-half outputs is +//! whole 8-column vectors). +constexpr int kDraftEhprojTableMaxRanks = 16; + +//! Whether invokeDraftEhprojTable serves tpSize ranks. +constexpr bool draftEhprojTableSupportsTpSize(int tpSize) +{ + return tpSize >= 2 && tpSize <= kDraftEhprojTableMaxRanks && kDraftSampleGlueHidden % (8 * tpSize) == 0; +} + +//! Bytes one rank publishes per row in invokeDraftEhprojTable: its (index, value) packet, the BF16 token-table row of +//! its candidate token and its kDraftSampleGlueHidden / tpSize FP32 hidden-half eh_proj outputs. +constexpr int draftEhprojTableSlotBytes(int tpSize) +{ + return kDraftSampleGluePacketBytes + 2 * kDraftSampleGlueHidden + 4 * kDraftSampleGlueHidden / tpSize; +} + +//! One MTP draft step after the vocabulary-sharded draft LM head that also replaces the next step's eh_proj and its +//! all-reduce + input_layernorm, in one cluster of kDraftEhprojTableClusterSize CTAs: +//! - the greedy draft token and the next step's state, as invokeDraftSampleGlue computes them; +//! - the next step's eh_proj output x = BF16(tokenTable[token] + hiddenHalf), one FP32 add, where tokenTable holds +//! eh_proj's embedding half applied to enorm(embed(t)) for the tokens t of this rank's vocabulary shard (BF16, one +//! rounding of the FP32 sum), and hiddenHalf is eh_proj's hidden half applied to hnorm(hidden) in FP32, the ranks +//! computing H / tpSize outputs each. Every rank publishes the token-table row of its candidate and its share of +//! hiddenHalf with its (index, value) packet in one exchange; every rank then takes the winner's row; +//! - input_layernorm(x) and x, what eh_proj's all-reduce with input_layernorm and a zero residual returns. +struct DraftEhprojTableParams +{ + //! Token, state and exchange workspace as for invokeDraftSampleGlue; glueOut and nextNewTokens must be null. The + //! workspace buffers hold numRows * tpSize slots of draftEhprojTableSlotBytes(tpSize) bytes. + DraftSampleGlueParams step; + //! This rank's hiddenHalf columns [rank * H / tpSize, (rank + 1) * H / tpSize), FP32 [numRows, H / tpSize] with + //! row stride hiddenHalfStride. + float const* hiddenHalf{nullptr}; + int64_t hiddenHalfStride{0}; + //! [vocabShard, H]: row t is eh_proj's embedding half applied to enorm(embed(rank * vocabShard + t)). + __nv_bfloat16 const* tokenTable{nullptr}; + //! input_layernorm weight [H] and epsilon. + __nv_bfloat16 const* normWeight{nullptr}; + float normEps{0.F}; + //! Outputs [numRows, H] with row stride outStride: input_layernorm(x) and x. + __nv_bfloat16* normedOut{nullptr}; + __nv_bfloat16* residualOut{nullptr}; + int64_t outStride{0}; +}; + +void invokeDraftEhprojTable(DraftEhprojTableParams const& params, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.cu new file mode 100644 index 000000000000..a73b31395b79 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.cu @@ -0,0 +1,1968 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.h" + +#include "tensorrt_llm/common/envUtils.h" + +#include + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::hist_topk +{ +namespace +{ + +constexpr int32_t kCtas = kCtasPerRow; +constexpr int32_t kThreads = 1024; +static_assert(kThreads / 32 == 32, "the block scans assume 32 warps"); +constexpr uint32_t kFull = 0xFFFFFFFFu; +constexpr int32_t kMaxGroups = kBins / kGroupBins; +constexpr int32_t kFineWords = kBins / 2; +static_assert(kHistWords == kFineWords + kMaxGroups, "histogram words: bin-count pairs, then group totals"); +constexpr uint32_t kEmpty = 0xFFFFFFFFu; +//! Exchanged bin counts saturate here; a group holds at most one saturated bin. +constexpr uint32_t kSat = 0xFFFFu; +static_assert(2 * int64_t{kSat} > kMaxShard, "two saturated bins would hold more keys than a shard"); +//! float4 of logits per thread (the keys are staged in shared memory). +constexpr int32_t kVecs = 3; +static_assert(kMaxShard == kCtas * kThreads * 4 * kVecs, "three float4 per thread at the largest shard"); +constexpr int32_t kKeys = 4 * kVecs; +constexpr int32_t kOutPerCta = kTopK / kCtas; +constexpr int32_t kCtasPerLaunch = kCtas * kMaxRows; +//! Boundary sub-bins of at most this many pairs are ranked directly. +constexpr int32_t kSmallRank = 64; +constexpr int64_t kHistRegionWords = int64_t{kMaxRows} * kMaxRanks * kHistWords; +constexpr int64_t kOutRegionWords = int64_t{kMaxRows} * kTopK; + +enum StateSlot : int32_t +{ + kStateCount = 0, // 64-bit: CTAs of exchange launches (each adds one once it read the count) + kStateError = 2, + kStateRecords = 4 +}; + +constexpr int32_t kRecordWords = 4; +static_assert(kStateRecords + kMaxRows * kRecordWords == kStateWords, "state layout"); + +enum Mode : int32_t +{ + kModeIdentity = 0, // no exchange for the row + kModeFull = 1, // the boundary bin is taken whole + kModeTies = 2, // the exact key: its ties are taken by count, the higher ranks first + kModeCand = 3, // the boundary pairs are exchanged and resolved on every rank + kModeNext = 4 // the boundary bin is refined with the next digit +}; + +enum DebugSlot : int32_t +{ + // %globaltimer stamps + kDbgEntry = 0, + kDbgWaited = 1, + kDbgLoaded = 2, + kDbgLocalHist = 3, + kDbgB0 = 4, + kDbgSent0 = 5, + kDbgSuffix = 6, + kDbgBucketed = 7, + kDbgCoarse = 8, + kDbgGroup = 9, + kDbgFine = 10, + kDbgCut = 11, + kDbgPeer = 12, + kDbgClassify = 13, + kDbgStaged = 14, + kDbgFinalSent = 15, + kDbgOutCopied = 16, + kDbgBndPolled = 17, + kDbgPass1 = 18, + kDbgThreshold = 19, + kDbgResolved = 20, + kDbgEnd = 21, + kDbgCleared = 22, + kDbgPreWaited = 24, + kDbgTokSent = 25, + kDbgClockEntry = 26, // clock64 (not %globaltimer) + kDbgClockEnd = 27, + // statistics + kDbgCutGroup = 32, + kDbgCutBin = 33, + kDbgOwnAbove = 34, + kDbgLevels = 35, + kDbgMode = 36, + kDbgBoundary = 37, + kDbgNeed = 38, + kDbgOutOff = 39, + kDbgBndOff = 40, + kDbgSelCount = 41, + kDbgBndCount = 42, + kDbgLocalCut = 43 +}; + +static_assert(kDbgLocalCut < kDebugSlots, "debug slots"); + +// Shared memory (32-bit words). Peers never read this CTA's shared memory: the cluster's exchanges are pushes +// (st.async) that complete on the receiver's barriers. +constexpr int32_t kSmemHist = 0; // [kBins] local histogram of the current level +constexpr int32_t kSmemSuff = kSmemHist + kBins; // [kBins + 4] its suffix sums +constexpr int32_t kSmemKeys = kSmemSuff + kBins + 4; // [kVecs][kThreads] uint4 keys +constexpr int32_t kSmemStageArea = kSmemKeys + kKeys * kThreads; // [kStageWords] buckets or stage (see below) +constexpr int32_t kStageWords = kKeys * kThreads / 2; +constexpr int32_t kSmemRHist = kSmemStageArea + kStageWords; // [kBins] bucket fill, then the resolve histogram +constexpr int32_t kSmemSlots = kSmemRHist + kBins; // [kCtas][kBins / kCtas] peers' counts of our bins +constexpr int32_t kSmemCoarse = kSmemSlots + kBins; // [kMaxRanks][kMaxGroups] +constexpr int32_t kSmemFine = kSmemCoarse + kMaxRanks * kMaxGroups; // [kMaxRanks][kGroupBins] +constexpr int32_t kSmemGsum = kSmemFine + kMaxRanks * kGroupBins; // [2][kMaxGroups] group totals (all, lower ranks) +constexpr int32_t kPreRow = kGroupBins + 4; // the group's kGroupBins + 1 suffix sums, padded to 16-byte pushes +constexpr int32_t kSmemPre = kSmemGsum + 2 * kMaxGroups; // [kLevels][kCtas][kPreRow] CTAs' suffix sums +constexpr int32_t kSmemScratch = kSmemPre + kLevels * kCtas * kPreRow; // [256] scan scratch +constexpr int32_t kSmemCut = kSmemScratch + 256; // [32] the cut of the current level +constexpr int32_t kSmemOffs = kSmemCut + 32; // [8] this CTA's offsets (warp 0) +constexpr int32_t kSmemMisc = kSmemOffs + 8; // [16] +constexpr int32_t kSmemSmall = kSmemMisc + 16; // [2 * kSmallRank] +constexpr int32_t kSmemBar = kSmemSmall + 2 * kSmallRank; // [2 * kLevels] uint64 barriers +constexpr int32_t kSmemDbg = kSmemBar + 4 * kLevels; // [kDebugSlots] int64 (debug only) +constexpr int32_t kSmemWords = kSmemDbg + 2 * kDebugSlots; +constexpr size_t kSmemBytes = size_t{kSmemWords} * sizeof(uint32_t); +// The stage area holds, for a cut on the first digit, the level-0 buckets: 16-bit key indices of this CTA (the key at +// index i is keys[i / 4] component i % 4, its token base + cbase + i); for a later cut, the selected tokens. +constexpr int32_t kSmemBucket = kSmemStageArea; // [kKeys * kThreads] uint16 +constexpr int32_t kSmemStage = kSmemStageArea; // [kTopK] selected tokens +static_assert(kTopK <= kStageWords && kKeys * kThreads <= 65536, "16-bit bucket indices, the stage in the buckets"); +static_assert(kSmemKeys % 4 == 0 && kSmemStage % 4 == 0 && kSmemSlots % 4 == 0 && kSmemFine % 4 == 0 + && kSmemPre % 4 == 0 && kSmemScratch % 4 == 0 && kSmemBar % 2 == 0 && kSmemDbg % 2 == 0, + "aligned arrays"); +// Below the 132 KB carveout: L1 keeps the register spills of 1024 threads. +static_assert(kSmemBytes <= 131 * 1024, "shared memory"); + +// kSmemCut slots. +enum CutSlot : int32_t +{ + kCutBin = 0, + kCutAboveAll = 1, // all ranks' keys above the cut bin (at this level, within the active prefix) + kCutBinAll = 2, // all ranks' keys in the cut bin + kCutAboveLow = 3, // the lower ranks' keys above the cut bin + kCutGroup = 4, // the group of the cut bin + kCutGroupAbove = 5, // all ranks' keys in the groups above it + kCutGroupAboveLow = 6, // the lower ranks' keys in the groups above it + kCutBndRankOff = 7, // the lower ranks' keys in the cut bin + kCutTiesOff = 8, // the exact key's ties the lower ranks take (when the cut is the exact key) + kCutTakeRank = 9, // the ties this rank takes + kCutRank = 16 // [kMaxRanks] every rank's count in the cut bin +}; + +// kSmemOffs slots. +enum OffsSlot : int32_t +{ + kOffsOut = 0, // this CTA's first word in the output stage + kOffsBnd = 1, // this CTA's first pair in the boundary stage + kOffsTake = 2, // the exact-key ties this CTA takes (kModeTies) + kOffsBndCta = 3, // this CTA's keys in the boundary bin + kOffsAboveCta = 4 // this CTA's keys above the boundary bin +}; + +enum MiscSlot : int32_t +{ + kMiscLevelBin = 0, // [kLevels] the cut bin of each level + kMiscSmallCount = 4, + kMiscBucketBin = 5, // the lowest bin of the level-0 buckets + kMiscRec = 6, // [3] the previous exchange launch's record of the row (levels, output words, boundary pairs) + kMiscLevelGroup = 9 // [kLevels] the group of each level's cut bin +}; + +__device__ __forceinline__ int32_t levelShift(int32_t level) +{ + return level == 0 ? 21 : (level == 1 ? 10 : 0); +} + +__device__ __forceinline__ int32_t levelBits(int32_t level) +{ + return level == 2 ? 10 : 11; +} + +__device__ __forceinline__ uint32_t orderedKey(float v) +{ + uint32_t const b = __float_as_uint(v); + uint32_t const k = (b & 0x80000000u) ? ~b : (b | 0x80000000u); + // NaN keys above +inf are clamped so that no key equals the empty word. + return k < kEmpty ? k : kEmpty - 1u; +} + +__device__ __forceinline__ uint32_t nibbleCount(uint32_t mask, int32_t k) +{ + return static_cast(__popc((mask >> (4 * k)) & 0xFu)); +} + +__device__ __forceinline__ uint4 ldVolatile4(uint32_t const* p) +{ + uint4 v; + asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(p) + : "memory"); + return v; +} + +__device__ __forceinline__ void st1(uint32_t* p, uint32_t v) +{ + asm volatile("st.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} + +__device__ __forceinline__ void st4(uint32_t* p, uint4 v) +{ + asm volatile("st.global.v4.u32 [%0], {%1, %2, %3, %4};" ::"l"(p), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w) + : "memory"); +} + +__device__ __forceinline__ uint64_t globalTimer() +{ + uint64_t t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + return t; +} + +// The debug slots of this CTA live in shared memory and are copied to HistTopKParams::debug at the exit: a stamp adds +// no global store (a later release would wait for it). +__device__ __forceinline__ int64_t* debugSlots() +{ + extern __shared__ __align__(16) uint32_t dynSmem[]; + return reinterpret_cast(dynSmem + kSmemDbg); +} + +__device__ __forceinline__ void stamp(bool dbg, int32_t slot) +{ + if (dbg && threadIdx.x == 0) + { + debugSlots()[slot] = static_cast(globalTimer()); + } +} + +__device__ __forceinline__ void putDebug(bool dbg, int32_t slot, int64_t v) +{ + if (dbg && threadIdx.x == 0) + { + debugSlots()[slot] = v; + } +} + +// Thread 0: this CTA's debug slots to HistTopKParams::debug. +__device__ __forceinline__ void flushDebug(int64_t* debug, int32_t row, int32_t c) +{ + if (debug != nullptr && threadIdx.x == 0) + { + int64_t* const g = debug + (int64_t{row} * kCtas + c) * kDebugSlots; + for (int32_t i = 0; i < kDebugSlots; ++i) + { + g[i] = debugSlots()[i]; + } + } +} + +__device__ __forceinline__ void clusterArrive() +{ + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); +} + +// An arrive without release semantics: the arriving thread does not wait for its own stores to be performed. +__device__ __forceinline__ void clusterArriveRelaxed() +{ + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); +} + +__device__ __forceinline__ void clusterWait() +{ + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); +} + +__device__ __forceinline__ void clusterSync() +{ + __syncwarp(); + clusterArrive(); + clusterWait(); +} + +__device__ __forceinline__ uint32_t smemAddr(void const* ptr) +{ + return static_cast(__cvta_generic_to_shared(ptr)); +} + +// The shared::cluster address of the same shared variable in CTA `rank` of the cluster. +__device__ __forceinline__ uint32_t peerAddr(uint32_t addr, int32_t rank) +{ + uint32_t r; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(r) : "r"(addr), "r"(rank)); + return r; +} + +__device__ __forceinline__ void barInit(uint32_t bar, uint32_t count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" ::"r"(bar), "r"(count) : "memory"); +} + +// The phase's one arrival, expecting `bytes` of asynchronous stores before it completes. +__device__ __forceinline__ void barArriveExpect(uint32_t bar, uint32_t bytes) +{ + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(bar), "r"(bytes) : "memory"); +} + +// Spins on the phase with test_wait (try_wait may suspend the thread for longer than the phase takes to complete). +__device__ __forceinline__ void barWait(uint32_t bar, uint32_t parity) +{ + asm volatile( + "{\n" + ".reg .pred done;\n" + "BAR_WAIT:\n" + "mbarrier.test_wait.parity.shared::cta.b64 done, [%0], %1;\n" + "@!done bra BAR_WAIT;\n" + "}\n" ::"r"(bar), + "r"(parity) + : "memory"); +} + +// Orders the barrier initializations before the cluster arrive that follows (peers then push to them). +__device__ __forceinline__ void fenceBarInit() +{ + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); +} + +// 16 bytes into a peer's shared memory, completing on the peer's barrier. +__device__ __forceinline__ void pushV4(uint32_t remote, uint4 v, uint32_t remoteBar) +{ + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.u32 [%0], {%1, %2, %3, %4}, [%5];" ::"r"(remote), + "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w), "r"(remoteBar) + : "memory"); +} + +__device__ __forceinline__ void st2(uint32_t* p, uint32_t a, uint32_t b) +{ + asm volatile("st.global.v2.u32 [%0], {%1, %2};" ::"l"(p), "r"(a), "r"(b) : "memory"); +} + +__device__ __forceinline__ int64_t histOff(int32_t buf, int32_t level, int32_t row, int32_t rank) +{ + return int64_t{buf} * kBufferWords + (int64_t{level} * kMaxRows + row) * kMaxRanks * kHistWords + + int64_t{rank} * kHistWords; +} + +__device__ __forceinline__ int64_t outOff(int32_t buf, int32_t row) +{ + return int64_t{buf} * kBufferWords + kLevels * kHistRegionWords + int64_t{row} * kTopK; +} + +__device__ __forceinline__ int64_t bndOff(int32_t buf, int32_t row) +{ + return int64_t{buf} * kBufferWords + kLevels * kHistRegionWords + kOutRegionWords + int64_t{row} * kBoundaryCap * 2; +} + +struct WaitCtx +{ + int32_t* state; + uint64_t timeoutNs; + int32_t rank; +}; + +// what: 1..3 the histograms of a level, 10 the output stage, 11 the boundary stage, 50 the previous launch's words +// before clearing; 200 + level / 300 + level: the ranks' histograms disagree; 400: the shard exceeds the logits. +__device__ __noinline__ void abortKernel(WaitCtx const& w, int32_t what) +{ + if (atomicExch(&w.state[kStateError], what) == 0) + { + printf("[glm_dsa_hist_topk] rank %d, row %d: %s (%d); aborting the kernel\n", w.rank, + static_cast(blockIdx.y), what < 200 ? "a wait for a peer exceeded the timeout" : "bad input", + what); + } + __threadfence_system(); + __trap(); +} + +struct PollClock +{ + uint64_t start = 0; + int32_t iter = 0; + + __device__ __forceinline__ void tick(WaitCtx const& w, int32_t what) + { + if ((++iter & 255) == 0) + { + uint64_t const now = globalTimer(); + if (start == 0) + { + start = now; + } + else if (now - start > w.timeoutNs) + { + abortKernel(w, what); + } + } + } +}; + +__device__ __forceinline__ bool anyEmpty(uint4 v, int32_t valid) +{ + return (valid > 0 && v.x == kEmpty) || (valid > 1 && v.y == kEmpty) || (valid > 2 && v.z == kEmpty) + || (valid > 3 && v.w == kEmpty); +} + +// The first `valid` words at p once a peer wrote them. +__device__ __forceinline__ uint4 pollVec(uint32_t const* p, int32_t valid, WaitCtx const& wc, int32_t what) +{ + uint4 v = ldVolatile4(p); + PollClock clk; + while (anyEmpty(v, valid)) + { + clk.tick(wc, what); + v = ldVolatile4(p); + } + return v; +} + +// Exclusive prefixes of N words in thread order and their block totals. One barrier; ws[0, 32 N) must not be +// rewritten before the next barrier. +template +__device__ __forceinline__ void blockExclusiveScanN( + uint32_t const (&v)[N], uint32_t* ws, uint32_t (&pre)[N], uint32_t (&tot)[N]) +{ + int32_t const lane = threadIdx.x & 31; + int32_t const warp = threadIdx.x >> 5; + uint32_t x[N]; +#pragma unroll + for (int32_t n = 0; n < N; ++n) + { + x[n] = v[n]; + } +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { +#pragma unroll + for (int32_t n = 0; n < N; ++n) + { + uint32_t const y = __shfl_up_sync(kFull, x[n], o); + if (lane >= o) + { + x[n] += y; + } + } + } + if (lane == 31) + { +#pragma unroll + for (int32_t n = 0; n < N; ++n) + { + ws[32 * n + warp] = x[n]; + } + } + __syncthreads(); + uint32_t w[N]; + uint32_t y[N]; +#pragma unroll + for (int32_t n = 0; n < N; ++n) + { + w[n] = ws[32 * n + lane]; + y[n] = w[n]; + } +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { +#pragma unroll + for (int32_t n = 0; n < N; ++n) + { + uint32_t const z = __shfl_up_sync(kFull, y[n], o); + if (lane >= o) + { + y[n] += z; + } + } + } +#pragma unroll + for (int32_t n = 0; n < N; ++n) + { + tot[n] = __shfl_sync(kFull, y[n], 31); + pre[n] = __shfl_sync(kFull, y[n] - w[n], warp) + x[n] - v[n]; + } +} + +// Exclusive suffix of v in thread order (the sum over the higher threads) and the block total. One barrier; ws[0, 32) +// as above. +__device__ __forceinline__ uint32_t blockExclusiveSuffix(uint32_t v, uint32_t* ws, uint32_t& total) +{ + int32_t const lane = threadIdx.x & 31; + int32_t const warp = threadIdx.x >> 5; + uint32_t x = v; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_down_sync(kFull, x, o); + if (lane + o < 32) + { + x += y; + } + } + if (lane == 0) + { + ws[warp] = x; + } + __syncthreads(); + uint32_t const w = ws[lane]; + uint32_t y = w; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const z = __shfl_down_sync(kFull, y, o); + if (lane + o < 32) + { + y += z; + } + } + total = __shfl_sync(kFull, y, 0); + return __shfl_sync(kFull, y - w, warp) + x - v; +} + +__device__ __forceinline__ uint32_t blockExclusiveSuffix(uint32_t v, uint32_t* ws) +{ + uint32_t total = 0; + return blockExclusiveSuffix(v, ws, total); +} + +// Position bases, in position order (float4 k major, then thread, then lane j: increasing column), of the keys +// flagged in mask a and in mask b (bit 4 k + j): the position of this thread's first flagged key of float4 k is +// ba[k] (bb[k]), the next ones follow. Also the block counts of both. One block scan (16-bit fields: a CTA flags at +// most 4 keys per thread and float4). +__device__ __forceinline__ void positionBases( + uint32_t ma, uint32_t mb, uint32_t* ws, uint32_t (&ba)[kVecs], uint32_t (&bb)[kVecs], uint32_t& na, uint32_t& nb) +{ + static_assert(kVecs == 3, "three float4 counts per mask"); + uint32_t const v[3] = {nibbleCount(ma, 0) | (nibbleCount(ma, 1) << 16), + nibbleCount(ma, 2) | (nibbleCount(mb, 0) << 16), nibbleCount(mb, 1) | (nibbleCount(mb, 2) << 16)}; + uint32_t pre[3]; + uint32_t tot[3]; + blockExclusiveScanN<3>(v, ws, pre, tot); + uint32_t const ta0 = tot[0] & 0xFFFFu; + uint32_t const ta1 = tot[0] >> 16; + uint32_t const ta2 = tot[1] & 0xFFFFu; + uint32_t const tb0 = tot[1] >> 16; + uint32_t const tb1 = tot[2] & 0xFFFFu; + uint32_t const tb2 = tot[2] >> 16; + ba[0] = pre[0] & 0xFFFFu; + ba[1] = ta0 + (pre[0] >> 16); + ba[2] = ta0 + ta1 + (pre[1] & 0xFFFFu); + bb[0] = pre[1] >> 16; + bb[1] = tb0 + (pre[2] & 0xFFFFu); + bb[2] = tb0 + tb1 + (pre[2] >> 16); + na = ta0 + ta1 + ta2; + nb = tb0 + tb1 + tb2; +} + +// The bin of a local histogram (nbins <= 2 * kThreads, two adjacent bins per thread) holding the need-th largest +// item counting from the top bin: (bin, count above, bin count) into cut[0, 3). Two barriers. +__device__ __forceinline__ void findLocalCut( + uint32_t const* h, int32_t nbins, uint32_t need, uint32_t* ws, uint32_t* cut) +{ + int32_t const t = threadIdx.x; + bool const has = 2 * t < nbins; + uint32_t const c0 = has ? h[2 * t] : 0u; + uint32_t const c1 = has ? h[2 * t + 1] : 0u; + uint32_t const above1 = blockExclusiveSuffix(c0 + c1, ws); + uint32_t const above0 = above1 + c1; + if (has) + { + if (above1 < need && need <= above1 + c1) + { + cut[0] = 2 * t + 1; + cut[1] = above1; + cut[2] = c1; + } + else if (above0 < need && need <= above0 + c0) + { + cut[0] = 2 * t; + cut[1] = above0; + cut[2] = c0; + } + } + __syncthreads(); +} + +// Every CTA of the row owns nbins / kCtas bins of the level's histogram. This CTA pushes its local counts of each +// peer's bins into that peer's slots (st.async, completing on the peer's barrier `bar`), waits for the peers' counts of +// its own bins, and stores the rank's counts of its bins through the multicast address: two per word, saturated at +// kSat, and its 32-bin group totals. The local histogram must be complete (a CTA barrier before). red: >= 16 + 128 +// words of 16-byte aligned scratch. One barrier; the multicast stores are issued by the last two warps. +__device__ __forceinline__ void pushAndSend( + uint32_t const* hist, int32_t nbins, uint32_t* slots, uint32_t bar, uint32_t* red, uint32_t* mc) +{ + int32_t const t = threadIdx.x; + int32_t const perCta = nbins / kCtas; // 256 or 128 bins: whole warps, whole groups + int32_t const groupsPerCta = perCta / kGroupBins; // 8 or 4 + int32_t const c = static_cast(blockIdx.x); + int32_t const vecs = perCta / 4; + for (int32_t i = t; i < kCtas * vecs; i += kThreads) + { + int32_t const q = i / vecs; + if (q != c) + { + int32_t const k = i - q * vecs; + pushV4(peerAddr(smemAddr(slots + c * perCta + 4 * k), q), + *reinterpret_cast(hist + q * perCta + 4 * k), peerAddr(bar, q)); + } + } + if (t < perCta) + { + barWait(bar, 0); + uint32_t s = hist[c * perCta + t]; +#pragma unroll + for (int32_t q = 0; q < kCtas; ++q) + { + s += q != c ? slots[q * perCta + t] : 0u; + } + uint32_t const group = __reduce_add_sync(kFull, s); + if ((t & 31) == 0) + { + red[t / 32] = group; + } + // Lane 8 m packs the bins 8 m .. 8 m + 7 of the warp: four words of (even, odd) bins. + uint32_t const w = min(s, kSat) | (min(__shfl_down_sync(kFull, s, 1), kSat) << 16); + uint32_t const w2 = __shfl_down_sync(kFull, w, 2); + uint32_t const w4 = __shfl_down_sync(kFull, w, 4); + uint32_t const w6 = __shfl_down_sync(kFull, w, 6); + if ((t & 7) == 0) + { + *reinterpret_cast(red + 16 + t / 2) = make_uint4(w, w2, w4, w6); + } + } + __syncthreads(); + int32_t const u = t - (kThreads - 64); + if (u >= 0 && u < perCta / 8) + { + st4(mc + (c * perCta) / 2 + 4 * u, *reinterpret_cast(red + 16 + 4 * u)); + } + else if (u >= 32 && 4 * (u - 32) < groupsPerCta) + { + st4(mc + kFineWords + c * groupsPerCta + 4 * (u - 32), *reinterpret_cast(red + 4 * (u - 32))); + } +} + +// Suffix sums of a local histogram: suff[b] = keys in bins >= b, suff[nbins] = 0 (read by the cluster's CTAs). With +// bucketBin: the highest bin b with suff[b] >= min(kTopK, suff[0]) into *bucketBin (nbins without keys). The cut of +// all ranks cannot be above it: this CTA has fewer than kTopK keys above that cut. +__device__ __forceinline__ void localSuffix( + uint32_t const* hist, int32_t nbins, uint32_t* suff, uint32_t* ws, uint32_t* bucketBin) +{ + int32_t const t = threadIdx.x; + bool const has = 2 * t < nbins; + uint32_t const c0 = has ? hist[2 * t] : 0u; + uint32_t const c1 = has ? hist[2 * t + 1] : 0u; + uint32_t total = 0; + uint32_t const above = blockExclusiveSuffix(c0 + c1, ws, total); + if (has) + { + suff[2 * t + 1] = above + c1; + suff[2 * t] = above + c1 + c0; + uint32_t const thr = min(static_cast(kTopK), total); + if (bucketBin != nullptr && thr > 0u && above < thr) + { + if (above + c1 >= thr) + { + *bucketBin = static_cast(2 * t + 1); + } + else if (above + c1 + c0 >= thr) + { + *bucketBin = static_cast(2 * t); + } + } + } + if (t == 0) + { + suff[nbins] = 0; + if (bucketBin != nullptr && total == 0u) + { + *bucketBin = static_cast(nbins); + } + } + __syncthreads(); +} + +// Clear, on this rank's copy, what the previous exchange launch wrote into buffer `buf` for this row. That launch read +// only the group totals and one group of each histogram, so every histogram word is first read written (no peer +// store of that launch can land after the clear), then cleared; the output and boundary stages were read whole. +__device__ __forceinline__ void clearPrevious(HistTopKParams const& p, int32_t buf, int32_t row, int32_t levels, + int32_t outWords, int32_t bndPairs, WaitCtx const& wc, int32_t ctaIndex, int32_t ctaCount) +{ + int32_t const c = ctaIndex; + int32_t const t = threadIdx.x; + uint4 const empty = make_uint4(kEmpty, kEmpty, kEmpty, kEmpty); + for (int32_t level = 0; level < levels; ++level) + { + // Written per rank at this level: nbins / 2 bin-count words and nbins / 32 group totals. + int32_t const nbins = 1 << levelBits(level); + int32_t const fineVecs = nbins / 8; + int32_t const rankVecs = fineVecs + nbins / (4 * kGroupBins); + uint32_t* base = p.ucWorkspace + histOff(buf, level, row, 0); + for (int32_t i = c * kThreads + t; i < p.worldSize * rankVecs; i += ctaCount * kThreads) + { + int32_t const r = i / rankVecs; + int32_t const j = i % rankVecs; + uint32_t* w = base + int64_t{r} * kHistWords + (j < fineVecs ? 4 * j : kFineWords + 4 * (j - fineVecs)); + (void) pollVec(w, 4, wc, 50); + st4(w, empty); + } + } + uint32_t* o = p.ucWorkspace + outOff(buf, row); + for (int32_t i = c * kThreads + t; i < outWords; i += ctaCount * kThreads) + { + st1(o + i, kEmpty); + } + uint32_t* b = p.ucWorkspace + bndOff(buf, row); + for (int32_t i = c * kThreads + t; i < 2 * bndPairs; i += ctaCount * kThreads) + { + st1(b + i, kEmpty); + } +} + +// The lowest set bit of a non-zero mask. +__device__ __forceinline__ int32_t lowBit(uint32_t m) +{ + return static_cast(__ffs(static_cast(m))) - 1; +} + +// Warp 0, after the final cut: every CTA's keys above it and in its bin (lane q < kCtas: CTA q's suffix sums pushed +// into pre at each level, waited for on preBar0 + 8 lv), then this CTA's offsets in the output and boundary stages +// (the rank's table from receiveCut, then the lower CTAs). Into offs[] (read after the next barrier). +__device__ __forceinline__ void cutOffsets(uint32_t const* misc, uint32_t const* pre, uint32_t preBar0, + uint32_t const* cut, int32_t level, int32_t mode, uint32_t lowTot, int32_t c, uint32_t* offs, bool dbg) +{ + int32_t const lane = static_cast(threadIdx.x & 31); + for (int32_t lv = 0; lv <= level; ++lv) + { + barWait(preBar0 + 8 * lv, 0); + } + stamp(dbg, kDbgPreWaited); + uint32_t aboveQ = 0; + uint32_t bndQ = 0; + if (lane < kCtas) + { + for (int32_t lv = 0; lv < level; ++lv) + { + uint32_t const j = misc[kMiscLevelBin + lv] - misc[kMiscLevelGroup + lv] * kGroupBins; + aboveQ += pre[(lv * kCtas + lane) * kPreRow + j + 1]; + } + uint32_t const* const pq = pre + (level * kCtas + lane) * kPreRow; + uint32_t const j = cut[kCutBin] - cut[kCutGroup] * kGroupBins; + aboveQ += pq[j + 1]; + bndQ = pq[j] - pq[j + 1]; + } + uint32_t rankOff = lowTot; + uint32_t bndRankOff = 0; + uint32_t takeQ = 0; + if (mode == kModeTies) + { + // The rank's ties go to the higher CTAs first (the larger tokens). + rankOff += cut[kCutTiesOff]; + int32_t const takeRank = static_cast(cut[kCutTakeRank]); + uint32_t xq = bndQ; +#pragma unroll + for (int32_t o = 1; o < kCtas; o <<= 1) + { + uint32_t const y = __shfl_down_sync(kFull, xq, o); + if (lane + o < kCtas) + { + xq += y; + } + } + int32_t const higherQ = static_cast(xq - bndQ); + takeQ = static_cast(max(0, min(takeRank - higherQ, static_cast(bndQ)))); + } + else + { + bndRankOff = cut[kCutBndRankOff]; + if (mode == kModeFull) + { + rankOff += bndRankOff; + takeQ = bndQ; + } + } + uint32_t const outOffC = rankOff + __reduce_add_sync(kFull, lane < c ? aboveQ + takeQ : 0u); + uint32_t const bndOffC = bndRankOff + __reduce_add_sync(kFull, lane < c ? bndQ : 0u); + uint32_t const takeC = __shfl_sync(kFull, takeQ, c); + uint32_t const bndC = __shfl_sync(kFull, bndQ, c); + uint32_t const aboveC = __shfl_sync(kFull, aboveQ, c); + if (lane == 0) + { + offs[kOffsOut] = outOffC; + offs[kOffsBnd] = bndOffC; + offs[kOffsTake] = takeC; + offs[kOffsBndCta] = bndC; + offs[kOffsAboveCta] = aboveC; + } +} + +// Store n words from shared memory through the multicast address: 16-byte stores where dst is aligned. +__device__ __forceinline__ void storeWords(uint32_t* dst, uint32_t const* src, int32_t n) +{ + int32_t const t = threadIdx.x; + int32_t const head + = min(n, static_cast((4u - static_cast((reinterpret_cast(dst) >> 2) & 3u)) & 3u)); + int32_t const vecs = (n - head) / 4; + int32_t const tail = n - head - 4 * vecs; + if (t < head) + { + st1(dst + t, src[t]); + } + for (int32_t i = t; i < vecs; i += kThreads) + { + int32_t const j = head + 4 * i; + st4(dst + j, make_uint4(src[j], src[j + 1], src[j + 2], src[j + 3])); + } + if (t >= 64 && t < 64 + tail) + { + int32_t const j = head + 4 * vecs + (t - 64); + st1(dst + j, src[j]); + } +} + +// Every CTA of the row: the cut of one level from all ranks' histograms, read through the group totals (4 KB) and the +// crossing group's bins (1 KB), with every rank's count in the cut bin and the rank's offsets that follow from them. +// Warps 2..8 meanwhile push this CTA's suffix sums over the crossing group's bins and the bin above them into row c +// of every CTA's pre rows of this level (completing on their barriers preBar0 + 8 level); when this is the final cut, +// warp 0 then computes this CTA's offsets (cutOffsets). gsum must be zero; lowTotPrev: the lower ranks' keys above +// the cuts of the earlier levels. Results in cut[] and offs[] (every thread can read them after the last barrier). +__device__ void receiveCut(HistTopKParams const& p, uint32_t const* ucLevel, uint32_t const* suffLevel, int32_t nbins, + int32_t level, uint32_t need, uint32_t lowTotPrev, uint32_t* coarse, uint32_t* fine, uint32_t* gsum, + uint32_t const* misc, uint32_t* pre, uint32_t preBar0, uint32_t* cut, uint32_t* offs, WaitCtx const& wc, + int32_t what, bool dbg) +{ + uint32_t* const preL = pre + level * kCtas * kPreRow; + uint32_t const preBar = preBar0 + 8 * level; + int32_t const t = threadIdx.x; + int32_t const world = p.worldSize; + int32_t const rank = p.rank; + int32_t const ngroups = nbins / kGroupBins; // 64 or 32 + int32_t const vecsPerRank = ngroups / 4; // 16 or 8 + uint32_t* const gAllSum = gsum; // [kMaxGroups] all ranks + uint32_t* const gLowSum = gsum + kMaxGroups; // [kMaxGroups] the lower ranks + if (t == 0) + { + cut[kCutGroup] = kEmpty; + cut[kCutBin] = kEmpty; + } + // 1. The group totals of every rank, summed per group over all ranks and over the lower ranks. + if (t < kMaxRanks * vecsPerRank) + { + int32_t const r = t / vecsPerRank; + int32_t const vi = t % vecsPerRank; + uint4 v = make_uint4(0, 0, 0, 0); + if (r < world) + { + v = pollVec(ucLevel + int64_t{r} * kHistWords + kFineWords + 4 * vi, 4, wc, what); + uint32_t const vv[4] = {v.x, v.y, v.z, v.w}; +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + atomicAdd(&gAllSum[4 * vi + j], vv[j]); + if (r < rank) + { + atomicAdd(&gLowSum[4 * vi + j], vv[j]); + } + } + } + coarse[r * kMaxGroups + 4 * vi] = v.x; + coarse[r * kMaxGroups + 4 * vi + 1] = v.y; + coarse[r * kMaxGroups + 4 * vi + 2] = v.z; + coarse[r * kMaxGroups + 4 * vi + 3] = v.w; + } + __syncthreads(); + stamp(dbg, kDbgCoarse); + // 2. The crossing group (warp 0): lane l holds groups 2 l and 2 l + 1. + if (t < 32) + { + uint32_t gAll[2]; + uint32_t gLow[2]; +#pragma unroll + for (int32_t k = 0; k < 2; ++k) + { + int32_t const g = 2 * t + k; + gAll[k] = g < ngroups ? gAllSum[g] : 0u; + gLow[k] = g < ngroups ? gLowSum[g] : 0u; + } + // Exclusive suffixes over the lanes of the per-lane totals. + uint32_t xa = gAll[0] + gAll[1]; + uint32_t xl = gLow[0] + gLow[1]; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const ya = __shfl_down_sync(kFull, xa, o); + uint32_t const yl = __shfl_down_sync(kFull, xl, o); + if (t + o < 32) + { + xa += ya; + xl += yl; + } + } + uint32_t const aAbove1 = xa - gAll[0] - gAll[1]; + uint32_t const lAbove1 = xl - gLow[0] - gLow[1]; + int32_t which = -1; + if (aAbove1 < need && need <= aAbove1 + gAll[1]) + { + which = 1; + } + else if (aAbove1 + gAll[1] < need && need <= aAbove1 + gAll[1] + gAll[0]) + { + which = 0; + } + if (which >= 0) + { + cut[kCutGroup] = 2 * t + which; + cut[kCutGroupAbove] = which == 1 ? aAbove1 : aAbove1 + gAll[1]; + cut[kCutGroupAboveLow] = which == 1 ? lAbove1 : lAbove1 + gLow[1]; + } + } + __syncthreads(); + int32_t const g = static_cast(cut[kCutGroup]); + if (g < 0 || g >= ngroups) + { + // No group holds the need-th key: the ranks' inputs disagree, nothing exact can follow. + abortKernel(wc, 200 + what); + } + uint32_t const needInGroup = need - cut[kCutGroupAbove]; + stamp(dbg, kDbgGroup); + // Warps 2..8 (idle in steps 3 and 4): this CTA's suffix sums over the group's bins and the bin above them + // (vectors 0..8 of four, the last three words padding) to every CTA of the row, one destination per warp: the 9 + // lanes' 16-byte pushes go to consecutive addresses of one peer, and the warps push at once. + if (t >= 64 && t < 64 + 32 * (kCtas - 1)) + { + int32_t const c = static_cast(blockIdx.x); + int32_t const lane = t & 31; + int32_t const qi = (t >> 5) - 2; + int32_t const q = qi < c ? qi : qi + 1; + if (lane < kPreRow / 4) + { + uint32_t const* const src = suffLevel + g * kGroupBins + 4 * lane; + uint4 const val = make_uint4(src[0], src[1], src[2], src[3]); + uint4* const slot = reinterpret_cast(preL + c * kPreRow) + lane; + if (qi == 0) + { + *slot = val; // this CTA's own row + } + pushV4(peerAddr(smemAddr(slot), q), val, peerAddr(preBar, q)); + } + } + // 3. The group's 32 bins of every rank (threads 4 r .. 4 r + 3: 8 bins each), a saturated bin restored from the + // group total. + if (t < kMaxRanks * 4) + { + int32_t const r = t / 4; + int32_t const vi = t % 4; + uint4 v = make_uint4(0, 0, 0, 0); + if (r < world) + { + v = pollVec(ucLevel + int64_t{r} * kHistWords + (kGroupBins / 2) * g + 4 * vi, 4, wc, what); + } + uint32_t const ww[4] = {v.x, v.y, v.z, v.w}; + uint32_t x[8]; + uint32_t others = 0; + bool sat = false; +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + x[2 * j] = ww[j] & 0xFFFFu; + x[2 * j + 1] = ww[j] >> 16; + } +#pragma unroll + for (int32_t j = 0; j < 8; ++j) + { + if (x[j] == kSat) + { + sat = true; + } + else + { + others += x[j]; + } + } + others += __shfl_xor_sync(kFull, others, 1); + others += __shfl_xor_sync(kFull, others, 2); + if (sat) + { + uint32_t const total = coarse[r * kMaxGroups + g]; +#pragma unroll + for (int32_t j = 0; j < 8; ++j) + { + if (x[j] == kSat) + { + x[j] = total - others; + } + } + } + *reinterpret_cast(fine + r * kGroupBins + 8 * vi) = make_uint4(x[0], x[1], x[2], x[3]); + *reinterpret_cast(fine + r * kGroupBins + 8 * vi + 4) = make_uint4(x[4], x[5], x[6], x[7]); + // Warps 0 and 1 only: warps 2..8 may still push. + asm volatile("bar.sync 1, 64;" ::: "memory"); + } + stamp(dbg, kDbgFine); + // 4. The cut bin (warp 0: lane = bin of the group), every rank's count in it and the rank's offsets. + if (t < 32) + { + uint32_t all = 0; + uint32_t low = 0; +#pragma unroll + for (int32_t r = 0; r < kMaxRanks; ++r) + { + uint32_t const x = fine[r * kGroupBins + t]; + all += x; + low += r < rank ? x : 0u; + } + // Exclusive suffix over the 32 bins of the group. + uint32_t xa = all; + uint32_t xl = low; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const ya = __shfl_down_sync(kFull, xa, o); + uint32_t const yl = __shfl_down_sync(kFull, xl, o); + if (t + o < 32) + { + xa += ya; + xl += yl; + } + } + uint32_t const aAbove = xa - all; + uint32_t const lAbove = xl - low; + uint32_t const holder = __ballot_sync(kFull, aAbove < needInGroup && needInGroup <= aAbove + all); + if (holder != 0u) + { + int32_t const w = lowBit(holder); + uint32_t const aAboveW = __shfl_sync(kFull, aAbove, w); + uint32_t const allW = __shfl_sync(kFull, all, w); + uint32_t const lAboveW = __shfl_sync(kFull, lAbove, w); + uint32_t const bc = t < kMaxRanks ? fine[t * kGroupBins + w] : 0u; + if (t < kMaxRanks) + { + cut[kCutRank + t] = bc; + } + uint32_t const bndRankOff = __reduce_add_sync(kFull, t < rank ? bc : 0u); + // When the bin is the exact key: rank r takes what the higher ranks leave of the bin's need. + uint32_t const needBin = needInGroup - aAboveW; + uint32_t xs = bc; +#pragma unroll + for (int32_t o = 1; o < kMaxRanks; o <<= 1) + { + uint32_t const y = __shfl_down_sync(kFull, xs, o); + if (t + o < kMaxRanks) + { + xs += y; + } + } + int32_t const higher = static_cast(xs - bc); + int32_t const take = max(0, min(static_cast(needBin) - higher, static_cast(bc))); + uint32_t const tiesOff = __reduce_add_sync(kFull, t < rank ? static_cast(take) : 0u); + uint32_t const takeRank = __shfl_sync(kFull, static_cast(take), rank); + uint32_t const aboveAll = cut[kCutGroupAbove] + aAboveW; + uint32_t const aboveLow = cut[kCutGroupAboveLow] + lAboveW; + if (t == 0) + { + cut[kCutBin] = static_cast(g * kGroupBins + w); + cut[kCutAboveAll] = aboveAll; + cut[kCutBinAll] = allW; + cut[kCutAboveLow] = aboveLow; + cut[kCutBndRankOff] = bndRankOff; + cut[kCutTiesOff] = tiesOff; + cut[kCutTakeRank] = takeRank; + } + __syncwarp(); + stamp(dbg, kDbgCut); + // When this is the final cut (as the main loop decides it), this CTA's offsets at once. + uint32_t const needAfter = need - aboveAll; + int32_t const mode = needAfter == allW + ? kModeFull + : (level == kLevels - 1 ? kModeTies + : (allW <= static_cast(kBoundaryCap) ? kModeCand : kModeNext)); + if (mode != kModeNext) + { + cutOffsets(misc, pre, preBar0, cut, level, mode, lowTotPrev + aboveLow, + static_cast(blockIdx.x), offs, dbg); + } + } + } + __syncthreads(); + stamp(dbg, kDbgPeer); +} + +// Store the tokens of n bucket entries (key indices) through the multicast address: 16-byte stores where dst is +// aligned. +__device__ __forceinline__ void storeBucketTokens(uint32_t* dst, uint16_t const* idx, uint32_t tokenBase, int32_t n) +{ + int32_t const t = threadIdx.x; + int32_t const head + = min(n, static_cast((4u - static_cast((reinterpret_cast(dst) >> 2) & 3u)) & 3u)); + int32_t const vecs = (n - head) / 4; + int32_t const tail = n - head - 4 * vecs; + if (t < head) + { + st1(dst + t, tokenBase + idx[t]); + } + for (int32_t i = t; i < vecs; i += kThreads) + { + int32_t const j = head + 4 * i; + st4(dst + j, + make_uint4(tokenBase + idx[j], tokenBase + idx[j + 1], tokenBase + idx[j + 2], tokenBase + idx[j + 3])); + } + if (t >= 64 && t < 64 + tail) + { + int32_t const j = head + 4 * vecs + (t - 64); + st1(dst + j, tokenBase + idx[j]); + } +} + +// Store n bucket entries as (key, token) pairs through the multicast address: 16-byte stores (two pairs) where dst +// is aligned (dst is 8-byte aligned). From the second half of the CTA, beside storeBucketTokens. +__device__ __forceinline__ void storeBucketPairs( + uint32_t* dst, uint16_t const* idx, uint32_t const* keys, uint32_t tokenBase, int32_t n) +{ + int32_t const t = (static_cast(threadIdx.x) + kThreads / 2) % kThreads; + int32_t const head = min(n, ((reinterpret_cast(dst) >> 3) & 1u) != 0u ? 1 : 0); + int32_t const vecs = (n - head) / 2; + if (t < head) + { + st1(dst, keys[idx[0]]); + st1(dst + 1, tokenBase + idx[0]); + } + for (int32_t i = t; i < vecs; i += kThreads) + { + int32_t const j = head + 2 * i; + uint32_t const a = idx[j]; + uint32_t const b = idx[j + 1]; + st4(dst + 2 * j, make_uint4(keys[a], tokenBase + a, keys[b], tokenBase + b)); + } + if (t == 64 && head + 2 * vecs < n) + { + uint32_t const a = idx[n - 1]; + st1(dst + 2 * (n - 1), keys[a]); + st1(dst + 2 * (n - 1) + 1, tokenBase + a); + } +} + +// The resolver: among all ranks' boundary pairs (thread t holds pairs 2 t, 2 t + 1, 2 kThreads + 2 t and +// 2 kThreads + 2 t + 1 in pk/pt, valid as flagged in pv; every key shares its bits from `shift` up with keyTop), +// select the `need` largest by (key, token) and write their tokens, in pair order, to outRow[base, base + need). +__device__ void resolveBoundary(uint32_t const (&pk)[4], uint32_t const (&pt)[4], uint32_t pv, uint32_t need, + int32_t shift, uint32_t keyTop, uint32_t* rh, uint32_t* ws, uint32_t* cut, uint32_t* small, uint32_t* smallCount, + int32_t* outRow, int32_t base, bool dbg) +{ + int32_t const t = threadIdx.x; + // The next 11 key bits (rh and smallCount are zero: cleared at the kernel start). + int32_t const w1 = shift > 11 ? 11 : shift; + int32_t const sh1 = shift - w1; + uint32_t const mask1 = (1u << w1) - 1u; +#pragma unroll + for (int32_t i = 0; i < 4; ++i) + { + if ((pv >> i) & 1u) + { + atomicAdd(&rh[(pk[i] >> sh1) & mask1], 1u); + } + } + __syncthreads(); + findLocalCut(rh, 1 << w1, need, ws, cut); + uint32_t const d1 = cut[0]; + uint32_t const need1 = need - cut[1]; + uint32_t const cnt1 = cut[2]; + stamp(dbg, kDbgPass1); + uint32_t tKey = 0; + uint32_t tTok = 0; + if (cnt1 <= static_cast(kSmallRank)) + { + // Few pairs in that sub-bin: rank them by (key, token) directly. +#pragma unroll + for (int32_t i = 0; i < 4; ++i) + { + if (((pv >> i) & 1u) && ((pk[i] >> sh1) & mask1) == d1) + { + uint32_t const slot = atomicAdd(smallCount, 1u); + small[slot] = pk[i]; + small[kSmallRank + slot] = pt[i]; + } + } + __syncthreads(); + if (t < static_cast(cnt1)) + { + uint32_t const k = small[t]; + uint32_t const tok = small[kSmallRank + t]; + uint32_t higher = 0; + for (uint32_t j = 0; j < cnt1; ++j) + { + uint32_t const kj = small[j]; + uint32_t const tj = small[kSmallRank + j]; + higher += (kj > k || (kj == k && tj > tok)) ? 1u : 0u; + } + if (higher == need1 - 1) + { + cut[0] = k; + cut[1] = tok; + } + } + __syncthreads(); + tKey = cut[0]; + tTok = cut[1]; + } + else + { + // The remaining key bits, then the token bits of the threshold key's ties (the larger tokens first). + uint32_t needR = need1; + uint32_t prefix = d1; + uint32_t tie = cnt1; + if (sh1 > 0) + { + __syncthreads(); + for (int32_t i = t; i < (1 << sh1); i += kThreads) + { + rh[i] = 0; + } + __syncthreads(); + uint32_t const mask2 = (1u << sh1) - 1u; +#pragma unroll + for (int32_t i = 0; i < 4; ++i) + { + if (((pv >> i) & 1u) && ((pk[i] >> sh1) & mask1) == d1) + { + atomicAdd(&rh[pk[i] & mask2], 1u); + } + } + __syncthreads(); + findLocalCut(rh, 1 << sh1, needR, ws, cut); + prefix = (d1 << sh1) | cut[0]; + needR -= cut[1]; + tie = cut[2]; + __syncthreads(); + } + uint32_t const kLow = prefix; + uint32_t const lowMask = (1u << shift) - 1u; + uint32_t tokPrefix = 0; + if (needR != tie) + { + int32_t rem = 21; + while (rem > 0) + { + int32_t const width = rem > 11 ? 11 : rem; + int32_t const sh = rem - width; + uint32_t const matchMask = (1u << (21 - rem)) - 1u; + for (int32_t i = t; i < (1 << width); i += kThreads) + { + rh[i] = 0; + } + __syncthreads(); +#pragma unroll + for (int32_t i = 0; i < 4; ++i) + { + if (((pv >> i) & 1u) && (pk[i] & lowMask) == kLow && ((pt[i] >> rem) & matchMask) == tokPrefix) + { + atomicAdd(&rh[(pt[i] >> sh) & ((1u << width) - 1u)], 1u); + } + } + __syncthreads(); + findLocalCut(rh, 1 << width, needR, ws, cut); + tokPrefix = (tokPrefix << width) | cut[0]; + needR -= cut[1]; + rem = sh; + __syncthreads(); + } + } + tKey = keyTop | kLow; + tTok = tokPrefix; + } + stamp(dbg, kDbgThreshold); + // The pairs at or above (tKey, tTok), in pair order. + uint32_t win = 0; +#pragma unroll + for (int32_t i = 0; i < 4; ++i) + { + if (((pv >> i) & 1u) && (pk[i] > tKey || (pk[i] == tKey && pt[i] >= tTok))) + { + win |= 1u << i; + } + } + uint32_t const v[1] = {static_cast(__popc(win & 3u)) | (static_cast(__popc(win >> 2)) << 16)}; + uint32_t pre[1]; + uint32_t tot[1]; + blockExclusiveScanN<1>(v, ws, pre, tot); + uint32_t p0 = pre[0] & 0xFFFFu; + uint32_t p1 = (tot[0] & 0xFFFFu) + (pre[0] >> 16); +#pragma unroll + for (int32_t i = 0; i < 4; ++i) + { + if ((win >> i) & 1u) + { + uint32_t& pos = i < 2 ? p0 : p1; + outRow[base + static_cast(pos)] = static_cast(pt[i]); + ++pos; + } + } +} + +__global__ void __launch_bounds__(kThreads, 1) dsaHistTopKKernel(HistTopKParams p) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + extern __shared__ __align__(16) uint32_t smem[]; + uint32_t* const hist = smem + kSmemHist; + uint32_t* const suff = smem + kSmemSuff; + uint4* const keys = reinterpret_cast(smem + kSmemKeys); + uint32_t* const stageTok = smem + kSmemStage; + uint16_t* const bucket = reinterpret_cast(smem + kSmemBucket); + uint32_t* const rhist = smem + kSmemRHist; + uint32_t* const cursor = rhist; // the level-0 bucket fill; its bins are cleared again for the resolver + uint32_t* const slots = smem + kSmemSlots; + uint32_t* const coarse = smem + kSmemCoarse; + uint32_t* const fine = smem + kSmemFine; + uint32_t* const gsum = smem + kSmemGsum; + uint32_t* const pre = smem + kSmemPre; + uint32_t* const ws = smem + kSmemScratch; + uint32_t* const cut = smem + kSmemCut; + uint32_t* const offs = smem + kSmemOffs; + uint32_t* const misc = smem + kSmemMisc; + uint32_t* const small = smem + kSmemSmall; + // Barriers: [kLevels] the peers' histogram counts, then [kLevels] the peers' suffix sums at the cut. + uint32_t const histBar0 = smemAddr(smem + kSmemBar); + uint32_t const preBar0 = histBar0 + 8 * kLevels; + + int32_t const c = static_cast(blockIdx.x); + int32_t const row = static_cast(blockIdx.y); + int32_t const t = threadIdx.x; + bool const dbg = p.debug != nullptr; + if (dbg && t == 0) + { + for (int32_t i = 0; i < kDebugSlots; ++i) + { + debugSlots()[i] = 0; + } + } + stamp(dbg, kDbgEntry); + putDebug(dbg, kDbgClockEntry, static_cast(clock64())); + + // The request's KV length is written before the logits of this call (the forward's metadata). + int32_t const kvLen = p.kvLens[0]; + int32_t const world = p.worldSize; + int32_t const rank = p.rank; + int32_t const perRank = (max(kvLen, 0) + world - 1) / world; + int32_t const S = (perRank + kShardAlign - 1) / kShardAlign * kShardAlign; + int32_t const base = rank * S; + int32_t const n = max(0, min(kvLen - base, S)); + int32_t const rowWin = kvLen - p.nextN + row + 1; + bool const rowActive = row < p.numRows; + bool const exchangeLaunch = kvLen - p.nextN + p.numRows > kTopK; + bool const rowExchange = rowActive && rowWin > kTopK; + int32_t const v = max(0, min(rowWin - base, n)); + int32_t const chunk = S / kCtas; + int32_t const cbase = c * chunk; + + for (int32_t i = t; i < kBins; i += kThreads) + { + hist[i] = 0; + rhist[i] = 0; + } + if (t == 0) + { + misc[kMiscSmallCount] = 0; + if (rowExchange) + { + // One arrival per phase (this thread, now), completing when the peers' bytes arrived. + for (int32_t lv = 0; lv < kLevels; ++lv) + { + barInit(histBar0 + 8 * lv, 1); + barInit(preBar0 + 8 * lv, 1); + barArriveExpect(histBar0 + 8 * lv, (kCtas - 1) * (1 << levelBits(lv)) / kCtas * 4); + barArriveExpect(preBar0 + 8 * lv, (kCtas - 1) * kPreRow * 4); + } + fenceBarInit(); + } + } + __syncthreads(); // the zeroed histograms before any thread counts into them, the barriers before their waits + if (rowExchange) + { + // Every CTA's barriers are initialized before any peer pushes to them (waited for after the logits load). + __syncwarp(); + clusterArriveRelaxed(); + } + + asm volatile("griddepcontrol.wait;" ::: "memory"); + stamp(dbg, kDbgWaited); + + int32_t* const outRow = p.out + int64_t{row} * kTopK; + if (rowActive && !rowExchange && t < kOutPerCta) + { + int32_t const pos = c * kOutPerCta + t; + outRow[pos] = pos < rowWin ? pos : -1; + } + if (!exchangeLaunch) + { + flushDebug(p.debug, row, c); + return; + } + WaitCtx const wc{p.state, p.timeoutNs, rank}; + if (S > p.localCapacity || S > kMaxShard) + { + abortKernel(wc, 400); + } + + // This CTA's chunk of the row's shard logits: keys into shared memory (each thread reads back only its own). + float const* rowLogits = p.logits + int64_t{row} * p.logitsStride + cbase; + float4 f[kVecs]; +#pragma unroll + for (int32_t k = 0; k < kVecs; ++k) + { + int32_t const q = t + k * kThreads; + f[k] = make_float4(0.F, 0.F, 0.F, 0.F); + if (rowExchange && 4 * q < chunk) + { + // Read once: not cached in L1, which then keeps the register spills. + f[k] = __ldcg(reinterpret_cast(rowLogits + 4 * q)); + } + } + // The epoch: every CTA of an exchange launch adds one to the count once it has read it (off the critical path: + // after its final send), and this launch started after the previous one completed, so a CTA reads base + k with + // base a multiple of kCtasPerLaunch and k < kCtasPerLaunch: count / kCtasPerLaunch is the number of earlier + // exchange launches. + unsigned long long* const counter = reinterpret_cast(p.state + kStateCount); + uint64_t const arrivals = *reinterpret_cast(counter); + int32_t const buf = static_cast((arrivals / kCtasPerLaunch) % kNumBuffers); + int32_t const prevBuf = (buf + kNumBuffers - 1) % kNumBuffers; + int32_t const* rec = p.state + kStateRecords + row * kRecordWords; + int32_t const recLevels = rec[0]; + int32_t const recOut = rec[1]; + int32_t const recBnd = rec[2]; + if (t == 0) + { + misc[kMiscRec] = static_cast(recLevels); + misc[kMiscRec + 1] = static_cast(recOut); + misc[kMiscRec + 2] = static_cast(recBnd); + } + if (rowExchange) + { + __syncwarp(); + clusterWait(); + } + + uint32_t validMask = 0; +#pragma unroll + for (int32_t k = 0; k < kVecs; ++k) + { + int32_t const q = t + k * kThreads; + if (rowExchange && 4 * q < chunk) + { + keys[q] = make_uint4(orderedKey(f[k].x), orderedKey(f[k].y), orderedKey(f[k].z), orderedKey(f[k].w)); + int32_t const col = cbase + 4 * q; +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + validMask |= (col + j < v ? 1u : 0u) << (4 * k + j); + } + } + } + stamp(dbg, kDbgLoaded); + + int32_t mode = kModeIdentity; + int32_t levelsUsed = 0; + uint32_t totalSel = 0; + uint32_t nbFinal = 0; + + if (!rowExchange) + { + clearPrevious(p, prevBuf, row, recLevels, recOut, recBnd, wc, c, kCtas); + if (t == 0) + { + atomicAdd(counter, 1ull); + } + clusterSync(); + } + else + { + uint32_t need = kTopK; + uint32_t aboveAllTot = 0; + uint32_t lowTot = 0; + uint32_t prefix = 0; + int32_t level = 0; + for (;; ++level) + { + int32_t const nbins = 1 << levelBits(level); + int32_t const shiftPrev = level == 0 ? 0 : levelShift(level - 1); + int32_t const shiftNow = levelShift(level); + uint32_t const maskNow = (1u << levelBits(level)) - 1u; +#pragma unroll + for (int32_t k = 0; k < kVecs; ++k) + { + uint32_t const vm = (validMask >> (4 * k)) & 0xFu; + if (vm != 0u) + { + uint4 const kv = keys[t + k * kThreads]; + uint32_t const kk[4] = {kv.x, kv.y, kv.z, kv.w}; +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + if (((vm >> j) & 1u) && (level == 0 || (kk[j] >> shiftPrev) == prefix)) + { + atomicAdd(&hist[(kk[j] >> shiftNow) & maskNow], 1u); + } + } + } + } + if (level == 0) + { + stamp(dbg, kDbgLocalHist); + } + __syncthreads(); // this CTA's local histogram of this level is complete + if (level == 0) + { + stamp(dbg, kDbgB0); + } + uint32_t* const ucLevel = p.ucWorkspace + histOff(buf, level, row, 0); + pushAndSend(hist, nbins, slots, histBar0 + 8 * level, ws, p.mcWorkspace + histOff(buf, level, row, rank)); + if (level == 0) + { + stamp(dbg, kDbgSent0); + } + if (t < 2 * kMaxGroups) + { + gsum[t] = 0; // receiveCut's sums, ordered before it by localSuffix's barrier + } + localSuffix(hist, nbins, suff, ws + 192, level == 0 ? misc + kMiscBucketBin : nullptr); + stamp(dbg, kDbgSuffix); + // The next level's local histogram starts from zero (receiveCut's barriers order this before it). + for (int32_t i = t; i < kBins; i += kThreads) + { + hist[i] = 0; + } + if (level == 0) + { + // This CTA's keys in the bins the cut cannot be above, by bin from the top (bin b at + // [suff[b + 1], suff[b])): the keys of a level-0 cut are then prefixes of the buckets. + uint32_t const bucketBin = misc[kMiscBucketBin]; +#pragma unroll + for (int32_t k = 0; k < kVecs; ++k) + { + uint32_t const vm = (validMask >> (4 * k)) & 0xFu; + if (vm != 0u) + { + uint4 const kv = keys[t + k * kThreads]; + uint32_t const kk[4] = {kv.x, kv.y, kv.z, kv.w}; +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + uint32_t const b = kk[j] >> shiftNow; + if (((vm >> j) & 1u) && b >= bucketBin) + { + uint32_t const pos = suff[b + 1] + atomicAdd(&cursor[b], 1u); + bucket[pos] = static_cast(4 * (t + k * kThreads) + j); + } + } + } + } + stamp(dbg, kDbgBucketed); + } + receiveCut(p, ucLevel, suff, nbins, level, need, lowTot, coarse, fine, gsum, misc, pre, preBar0, cut, offs, + wc, 1 + level, dbg); + uint32_t const bin = cut[kCutBin]; + if (bin == kEmpty) + { + abortKernel(wc, 300 + level); + } + putDebug(dbg, kDbgCutGroup, cut[kCutGroup]); + putDebug(dbg, kDbgCutBin, bin); + uint32_t const aboveAll = cut[kCutAboveAll]; + uint32_t const binAll = cut[kCutBinAll]; + uint32_t const aboveLow = cut[kCutAboveLow]; + if (t == 0) + { + misc[kMiscLevelBin + level] = bin; + misc[kMiscLevelGroup + level] = cut[kCutGroup]; + } + aboveAllTot += aboveAll; + lowTot += aboveLow; + need -= aboveAll; + prefix = level == 0 ? bin : ((prefix << levelBits(level)) | bin); + if (need == binAll) + { + mode = kModeFull; + } + else if (level == kLevels - 1) + { + mode = kModeTies; + } + else if (binAll <= static_cast(kBoundaryCap)) + { + mode = kModeCand; + } + else + { + mode = kModeNext; + } + if (mode != kModeNext) + { + break; + } + } + levelsUsed = level + 1; + nbFinal = cut[kCutBinAll]; + putDebug(dbg, kDbgLevels, levelsUsed); + putDebug(dbg, kDbgMode, mode); + putDebug(dbg, kDbgBoundary, nbFinal); + putDebug(dbg, kDbgNeed, need); + + // offs[]: this CTA's offsets, computed by warp 0 in receiveCut at the final cut. + + totalSel = mode == kModeCand ? aboveAllTot : static_cast(kTopK); + int32_t const shiftF = levelShift(level); + uint32_t selCount = 0; + uint32_t bndCount = 0; + uint32_t const outOffC = offs[kOffsOut]; + uint32_t const bndOffC = offs[kOffsBnd]; + uint32_t* const mcOut = p.mcWorkspace + outOff(buf, row) + outOffC; + uint32_t* const mcBnd = p.mcWorkspace + bndOff(buf, row) + 2 * int64_t{bndOffC}; + if (level == 0) + { + // The selected tokens and the boundary pairs are prefixes of the buckets (kModeFull or kModeCand). + uint32_t const nAbove = offs[kOffsAboveCta]; + bndCount = offs[kOffsBndCta]; + selCount = mode == kModeFull ? nAbove + bndCount : nAbove; + stamp(dbg, kDbgStaged); + uint32_t const tokenBase = static_cast(base + cbase); + storeBucketTokens(mcOut, bucket, tokenBase, static_cast(selCount)); + stamp(dbg, kDbgTokSent); + if (mode == kModeCand) + { + storeBucketPairs(mcBnd, bucket + nAbove, reinterpret_cast(keys), tokenBase, + static_cast(bndCount)); + } + } + else + { + // This CTA's keys above the final cut and in its bin. + uint32_t aboveMask = 0; + uint32_t bndMask = 0; +#pragma unroll + for (int32_t k = 0; k < kVecs; ++k) + { + uint32_t const vm = (validMask >> (4 * k)) & 0xFu; + if (vm != 0u) + { + uint4 const kv = keys[t + k * kThreads]; + uint32_t const kk[4] = {kv.x, kv.y, kv.z, kv.w}; +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + uint32_t const s = kk[j] >> shiftF; + if ((vm >> j) & 1u) + { + aboveMask |= (s > prefix ? 1u : 0u) << (4 * k + j); + bndMask |= (s == prefix ? 1u : 0u) << (4 * k + j); + } + } + } + } + stamp(dbg, kDbgClassify); + uint32_t bndBase[kVecs]; + uint32_t selBase[kVecs]; + uint32_t selMask = aboveMask; + if (mode != kModeTies) + { + if (mode == kModeFull) + { + selMask |= bndMask; + } + positionBases(bndMask, selMask, ws, bndBase, selBase, bndCount, selCount); + } + else + { + // The CTA's last offs[kOffsTake] keys of the exact key: the larger tokens. + uint32_t unusedBase[kVecs]; + uint32_t unusedCount = 0; + positionBases(bndMask, 0u, ws, bndBase, unusedBase, bndCount, unusedCount); + uint32_t const firstTaken = offs[kOffsBndCta] - offs[kOffsTake]; +#pragma unroll + for (int32_t k = 0; k < kVecs; ++k) + { + uint32_t pos = bndBase[k]; +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + uint32_t const bit = 1u << (4 * k + j); + if (bndMask & bit) + { + selMask |= pos >= firstTaken ? bit : 0u; + ++pos; + } + } + } + positionBases(selMask, 0u, ws + 128, selBase, unusedBase, selCount, unusedCount); + } + stamp(dbg, kDbgStaged); + + // The selected tokens into every rank's output stage, staged here so that the multicast stores are 16 + // bytes wide; in kModeCand the boundary pairs straight into every rank's boundary stage (8 bytes each). + uint32_t const pairMask = mode == kModeCand ? bndMask : 0u; +#pragma unroll + for (int32_t k = 0; k < kVecs; ++k) + { + if ((((selMask | pairMask) >> (4 * k)) & 0xFu) != 0u) + { + uint4 const kv = keys[t + k * kThreads]; + uint32_t const kk[4] = {kv.x, kv.y, kv.z, kv.w}; + uint32_t sp = selBase[k]; + uint32_t bp = bndBase[k]; +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + uint32_t const bit = 1u << (4 * k + j); + uint32_t const token = static_cast(base + cbase + 4 * (t + k * kThreads) + j); + if (selMask & bit) + { + stageTok[sp] = token; + ++sp; + } + else if (pairMask & bit) + { + st2(mcBnd + 2 * bp, kk[j], token); + } + bp += (bndMask & bit) ? 1u : 0u; + } + } + } + __syncthreads(); + storeWords(mcOut, stageTok, static_cast(selCount)); + } + stamp(dbg, kDbgFinalSent); + if (t == 0) + { + atomicAdd(counter, 1ull); // the epoch (see above), while the stage words travel + } + putDebug(dbg, kDbgSelCount, selCount); + putDebug(dbg, kDbgBndCount, bndCount); + putDebug(dbg, kDbgOutOff, outOffC); + putDebug(dbg, kDbgBndOff, bndOffC); + putDebug(dbg, kDbgOwnAbove, offs[kOffsAboveCta]); + + // The resolver requests all ranks' boundary pairs (nb <= kBoundaryCap: at most two 16-byte words per + // thread) before the output stage is copied. + int32_t const nb = static_cast(nbFinal); + bool const resolver = mode == kModeCand && c == kCtas - 1; + uint32_t const* ucBnd = p.ucWorkspace + bndOff(buf, row); + uint32_t const* ucOut = p.ucWorkspace + outOff(buf, row); + uint4 bv[2] = {make_uint4(0, 0, 0, 0), make_uint4(0, 0, 0, 0)}; + int32_t bvalid[2] = {0, 0}; + if (resolver) + { + int32_t const vecs = (2 * nb + 3) / 4; +#pragma unroll + for (int32_t k = 0; k < 2; ++k) + { + int32_t const vi = t + k * kThreads; + bvalid[k] = vi < vecs ? min(4, 2 * nb - 4 * vi) : 0; + if (bvalid[k] > 0) + { + bv[k] = ldVolatile4(ucBnd + 4 * vi); + } + } + // The level-0 bucket fill counters back to zero: they are the resolve histogram. + for (int32_t i = static_cast(misc[kMiscBucketBin]) + t; i < kBins; i += kThreads) + { + cursor[i] = 0; + } + __syncthreads(); + } + else + { + // What the previous exchange launch wrote into buffer prevBuf for this row, cleared by the CTAs other + // than the resolver while the peers' stage words travel. + clearPrevious(p, prevBuf, row, static_cast(misc[kMiscRec]), + static_cast(misc[kMiscRec + 1]), static_cast(misc[kMiscRec + 2]), wc, c, + mode == kModeCand ? kCtas - 1 : kCtas); + stamp(dbg, kDbgCleared); + } + // The output stage of every rank into the output row: this CTA's kOutPerCta positions. + if (t < kOutPerCta / 4) + { + int32_t const pos = c * kOutPerCta + 4 * t; + int32_t const valid = min(4, max(0, static_cast(totalSel) - pos)); + if (valid > 0) + { + uint4 const w = pollVec(ucOut + pos, valid, wc, 10); + if (valid == 4) + { + *reinterpret_cast(outRow + pos) = make_int4(static_cast(w.x), + static_cast(w.y), static_cast(w.z), static_cast(w.w)); + } + else + { + uint32_t const ww[4] = {w.x, w.y, w.z, w.w}; + for (int32_t j = 0; j < valid; ++j) + { + outRow[pos + j] = static_cast(ww[j]); + } + } + } + } + stamp(dbg, kDbgOutCopied); + if (resolver) + { + PollClock clk; + while (anyEmpty(bv[0], bvalid[0]) || anyEmpty(bv[1], bvalid[1])) + { + clk.tick(wc, 11); +#pragma unroll + for (int32_t k = 0; k < 2; ++k) + { + if (anyEmpty(bv[k], bvalid[k])) + { + bv[k] = ldVolatile4(ucBnd + 4 * (t + k * kThreads)); + } + } + } + stamp(dbg, kDbgBndPolled); + uint32_t const pk[4] = {bv[0].x, bv[0].z, bv[1].x, bv[1].z}; + uint32_t const pt[4] = {bv[0].y, bv[0].w, bv[1].y, bv[1].w}; + uint32_t const pv = (bvalid[0] >= 2 ? 1u : 0u) | (bvalid[0] == 4 ? 2u : 0u) | (bvalid[1] >= 2 ? 4u : 0u) + | (bvalid[1] == 4 ? 8u : 0u); + resolveBoundary(pk, pt, pv, need, shiftF, prefix << shiftF, rhist, ws + 192, cut, small, + misc + kMiscSmallCount, outRow, static_cast(totalSel), dbg); + stamp(dbg, kDbgResolved); + } + } + + // This launch's record for the next exchange launch's clearing (every CTA of the row read the previous record + // before it pushed its level-0 counts, which CTA 0 waited for; a row without an exchange passed its cluster + // barrier). + if (t == 0 && c == 0) + { + int32_t* recW = p.state + kStateRecords + row * kRecordWords; + recW[0] = levelsUsed; + recW[1] = static_cast(totalSel); + recW[2] = mode == kModeCand ? static_cast(nbFinal) : 0; + } + stamp(dbg, kDbgEnd); + putDebug(dbg, kDbgClockEnd, static_cast(clock64())); + flushDebug(p.debug, row, c); +#else + (void) p; +#endif +} + +__global__ void dsaHistShardMetaKernel(int32_t const* kvLens, int32_t const* blockTable, int32_t blockTableWidth, + int32_t worldSize, int32_t rank, int32_t tokensPerBlock, int32_t* shardCtx, int32_t* shardBlockTable, + int32_t shardBlockTableWidth) +{ + int32_t const kvLen = kvLens[0]; + int32_t const perRank = (max(kvLen, 0) + worldSize - 1) / worldSize; + int32_t const size = (perRank + kShardAlign - 1) / kShardAlign * kShardAlign; + int32_t const base = rank * size; + int32_t const n = min(max(kvLen - base, 0), size); + if (threadIdx.x == 0) + { + shardCtx[0] = max(n, 1); + } + int32_t const firstPage = base / tokensPerBlock; + int32_t const pages = (n + tokensPerBlock - 1) / tokensPerBlock; + for (int32_t j = threadIdx.x; j < shardBlockTableWidth; j += blockDim.x) + { + shardBlockTable[j] = (j < pages && firstPage + j < blockTableWidth) ? blockTable[firstPage + j] : 0; + } +} + +int64_t defaultTimeoutNs() +{ + return kDefaultTimeoutMs * 1000000; +} + +} // namespace + +void invokeDsaHistTopK(HistTopKParams const& inParams, cudaStream_t stream) +{ + HistTopKParams params = inParams; + if (params.numRows <= 0 || params.numRows > kMaxRows || params.worldSize <= 0 || params.worldSize > kMaxRanks + || params.rank < 0 || params.rank >= params.worldSize || params.localCapacity > kMaxShard + || params.logitsStride % 4 != 0 || reinterpret_cast(params.logits) % 16 != 0 + || reinterpret_cast(params.state) % 8 != 0) + { + throw std::invalid_argument("glm_dsa_hist_topk: unsupported rows, group, shard capacity or alignment"); + } + if (params.timeoutNs == 0) + { + params.timeoutNs = static_cast(defaultTimeoutNs()); + } + static thread_local bool configured[64] = {}; + int device = 0; + cudaGetDevice(&device); + cudaError_t err = cudaSuccess; + if (device >= 64 || !configured[device]) + { + err = cudaFuncSetAttribute( + dsaHistTopKKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(kSmemBytes)); + if (err != cudaSuccess) + { + throw std::runtime_error(std::string("glm_dsa_hist_topk: smem attribute: ") + cudaGetErrorString(err)); + } + if (device < 64) + { + configured[device] = true; + } + } + cudaLaunchConfig_t config{}; + config.gridDim = dim3(kCtas, kMaxRows); + config.blockDim = dim3(kThreads); + config.dynamicSmemBytes = kSmemBytes; + config.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kCtas; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL() ? 1 : 0; + config.attrs = attrs; + config.numAttrs = 2; + err = cudaLaunchKernelEx(&config, dsaHistTopKKernel, params); + if (err == cudaSuccess) + { + err = cudaGetLastError(); + } + if (err != cudaSuccess) + { + throw std::runtime_error(std::string("glm_dsa_hist_topk: launch: ") + cudaGetErrorString(err)); + } +} + +void invokeDsaHistShardMeta(int32_t const* kvLens, int32_t const* blockTable, int32_t blockTableWidth, + int32_t worldSize, int32_t rank, int32_t tokensPerBlock, int32_t* shardCtx, int32_t* shardBlockTable, + int32_t shardBlockTableWidth, cudaStream_t stream) +{ + if (worldSize <= 0 || worldSize > kMaxRanks || rank < 0 || rank >= worldSize || tokensPerBlock <= 0 + || kShardAlign % tokensPerBlock != 0) + { + throw std::invalid_argument("glm_dsa_hist_shard_meta: unsupported group or page size"); + } + dsaHistShardMetaKernel<<<1, 256, 0, stream>>>(kvLens, blockTable, blockTableWidth, worldSize, rank, tokensPerBlock, + shardCtx, shardBlockTable, shardBlockTableWidth); + cudaError_t const err = cudaGetLastError(); + if (err != cudaSuccess) + { + throw std::runtime_error(std::string("glm_dsa_hist_shard_meta: launch: ") + cudaGetErrorString(err)); + } +} + +} // namespace kernels::glm_moe_dsa::hist_topk + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.h new file mode 100644 index 000000000000..a50ddc363457 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.h @@ -0,0 +1,142 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::hist_topk +{ + +// Key-sharded DSA indexer top-k of one generation request, exact, by histogram exchange. +// +// The indexer context of the request is split over the ranks of a tensor-parallel group into contiguous shards: +// S = ceil(ceil(kvLen / worldSize) / kShardAlign) * kShardAlign, rank r holds the tokens [r S, r S + n_r) with +// n_r = clamp(kvLen - r S, 0, S), and scores only those (the caller supplies the shard's logits). Verify row i of +// nextN sees the tokens [0, kvLen - nextN + i + 1). +// +// Per row, one cluster of kCtasPerRow CTAs holds the row's shard keys (order-preserving keys of the FP32 logits). +// Every rank +// 1. builds the histogram of its keys' top 11 bits (per CTA in shared memory, reduced over the cluster: every CTA +// owns 1/kCtasPerRow of the bins and receives the other CTAs' counts of them as asynchronous stores into its +// shared memory, which complete on its barrier) and stores it once through the multicast address: the 32-bit +// totals of its 32-bin groups, and two 16-bit bin counts per word, saturated at 0xFFFF (at most one bin of a +// group can saturate, as two would hold more than kMaxShard keys; the receivers recover its count from the group +// total); +// 2. reads all ranks' group totals, finds the group that holds the kTopK-th largest key, reads that group's bins +// of all ranks and finds the cut bin: the keys above it are selected, and every rank knows each rank's count +// above it and in it; +// 3. when that bin is not taken whole and holds more than kBoundaryCap keys, refines the same way with the next +// 11 bits of the keys in the bin, then with the last 10 bits (the exact key, whose ties are taken from the +// higher ranks first: the larger tokens first); +// 4. stores its selected tokens through the multicast address into every rank's output stage, at its offset (rank, +// then CTA), and, when the boundary bin is neither taken whole nor the exact key, its boundary (key, token) pairs +// into every rank's boundary stage. After a cut on the top 11 bits (the common case) a CTA's tokens and pairs are +// prefixes of buckets it fills while it waits for the peers' histograms: its keys in the bins that cut cannot be +// above, by bin from the top, in the order of shared-memory atomics within a bin; after a later cut they are in +// position order; +// 5. copies the output stage into the output row; one CTA selects the boundary pairs that complete the row, by key +// and then the larger token, and appends them in stage order. +// Every rank derives each decision from the same multicast data, so all ranks write the same bits: the exact top-kTopK +// of the row by (logit, token), the larger token first among equal logits. The order within the row is the stages' +// order: the same on all ranks, but after a cut on the top 11 bits it can differ between calls with the same logits. +// Rows of at most kTopK tokens get [0, n) and -1 without an exchange, and a launch whose rows all have at most kTopK +// tokens touches neither the workspace nor the state. +// +// Workspace: kNumBuffers Lamport buffers, rotated by a device epoch that counts exchange launches. In exchange launch x +// a rank clears, on its own copy, what launch x - 1 wrote into buffer (x - 1) % kNumBuffers (the sizes are recorded +// in the state per row), each histogram word after it reads the word written (launch x - 1 read only the group +// totals and one group's bins); a peer writes that buffer again only in launch x + 2, after this rank completed +// launch x + 1. +// +// A wait for a peer's word that exceeds the timeout records the wait in the state's error word and aborts the kernel: +// the selection can no longer complete, and a late peer write could land in a buffer this rank reuses. + +constexpr int32_t kTopK = 2048; +constexpr int32_t kMaxRanks = 16; +constexpr int32_t kMaxRows = 8; +//! CTAs per row (one cluster). +constexpr int32_t kCtasPerRow = 8; +constexpr int32_t kShardAlign = 128; +//! Largest shard (tokens per rank): three float4 per thread of kCtasPerRow CTAs of 1024 threads. At 16 ranks a request +//! of up to 1,572,864 indexer tokens. +constexpr int32_t kMaxShard = 98304; +//! Boundary-bin keys exchanged per row; a fuller boundary bin is refined with the next digit instead. +constexpr int32_t kBoundaryCap = 4096; +constexpr int32_t kNumBuffers = 3; +//! Bins of the first two digits (11 bits); the last digit has 10 bits. +constexpr int32_t kBins = 2048; +//! Bins per group of the coarse summary. +constexpr int32_t kGroupBins = 32; +//! Lamport words of one (level, row, rank) histogram: kBins / 2 words of two 16-bit bin counts, then the 32-bit totals +//! of the kBins / kGroupBins groups. +constexpr int32_t kHistWords = kBins / 2 + kBins / kGroupBins; +constexpr int32_t kLevels = 3; +constexpr int64_t kBufferWords = int64_t{kLevels} * kMaxRows * kMaxRanks * kHistWords + int64_t{kMaxRows} * kTopK + + int64_t{kMaxRows} * kBoundaryCap * 2; +//! 32-bit words of the exchange workspace; every word must start as 0xFFFFFFFF. +constexpr int64_t kWorkspaceWords = kNumBuffers * kBufferWords; +//! int32 words of the per-rank device state (zero-initialized): a 64-bit count of the CTAs that completed exchange +//! launches (the epoch is that count / (kCtasPerRow * kMaxRows)), the error word, a spare word, then per row the +//! levels, output-stage words and boundary pairs the last exchange launch wrote. +constexpr int32_t kStateWords = 4 + kMaxRows * 4; +//! Wait for a peer's words before the kernel aborts, when HistTopKParams::timeoutNs is 0. +constexpr int64_t kDefaultTimeoutMs = 120000; +//! Debug slots per CTA (phase %globaltimer stamps and per-row statistics), for diagnosis only. +constexpr int32_t kDebugSlots = 64; + +struct HistTopKParams +{ + //! [numRows, logitsStride] FP32 logits of this rank's shard (columns [0, localCapacity), localCapacity >= S). + float const* logits; + int64_t logitsStride; + int32_t localCapacity; + //! [1] int32 indexer KV length of the request. + int32_t const* kvLens; + //! [numRows, kTopK] int32 output: request-local token indices, -1 for unused slots. + int32_t* out; + int32_t numRows; + int32_t nextN; + int32_t rank; + int32_t worldSize; + //! Multicast and this rank's unicast address of the exchange workspace (kWorkspaceWords words). + uint32_t* mcWorkspace; + uint32_t* ucWorkspace; + //! [kStateWords] int32 per-rank device state, 8-byte aligned. + int32_t* state; + //! Optional [kMaxRows, kCtasPerRow, kDebugSlots] int64 (nullptr = off). + int64_t* debug = nullptr; + //! Longest wait for a peer's words (ns); 0 = the default. + uint64_t timeoutNs = 0; +}; + +void invokeDsaHistTopK(HistTopKParams const& params, cudaStream_t stream); + +//! Shard of the generation request for `rank`: shardCtx[0] = max(n_r, 1) (an empty shard scores one token no row +//! reaches) and shardBlockTable[j] = blockTable[r S / tokensPerBlock + j] for the shard's pages, 0 beyond. One CTA. +void invokeDsaHistShardMeta(int32_t const* kvLens, int32_t const* blockTable, int32_t blockTableWidth, + int32_t worldSize, int32_t rank, int32_t tokensPerBlock, int32_t* shardCtx, int32_t* shardBlockTable, + int32_t shardBlockTableWidth, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa::hist_topk + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.cu new file mode 100644 index 000000000000..5224fcb38651 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.cu @@ -0,0 +1,467 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaDriverWrapper.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +extern "C" +{ + __device__ uint32_t __nvvm_get_smem_pointer(void*); +} + +namespace +{ + +using bf16_t = __nv_bfloat16; + +// Warp w issues the tensor copies of its two register stages w * kFusedAMma8PaceNs after the copy of its shared-memory +// stage (1.75 us over the CTA's eight warps); the CTA's share of the L2 prefetch ranges follows once its staged tiles +// are in registers. +constexpr int kFusedAMma8PaceNs = 250; + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 +__device__ void initialize_barrier(uint64_t* smem_barrier, // 64 bits user-manged barrier in smem + int thread_count = 1) // Thread count expected to arrive/wait on this barrier +{ + uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier); + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(smem_int_ptr), "r"(thread_count)); +} + +// Barrier wait +__device__ void wait_barrier(uint64_t* smem_barrier, // 64 bits user-manged barrier in smem + int phase_bit) // Current phase bit the barrier waiting to flip +{ + uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier); + asm volatile( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra DONE;\n" + "bra LAB_WAIT;\n" + "DONE:\n" + "}\n" ::"r"(smem_int_ptr), + "r"(phase_bit)); +} + +// Prefetches the ranges into L2 in 16 KB pieces. Piece q of the concatenated ranges goes to CTA q % gridDim.x: an SM +// issues bulk prefetches at a limited rate, so every CTA takes an equal share. +__device__ void prefetch_ranges_l2(FusedAL2Prefetch const& ranges, int local_tid, int thread_cnt) +{ + constexpr size_t piece_bytes = 16384; + size_t piece_cnt[FusedAL2Prefetch::kMaxRanges]; + size_t total_piece_cnt = 0; +#pragma unroll + for (int i = 0; i < FusedAL2Prefetch::kMaxRanges; i++) + { + piece_cnt[i] = (ranges.bytes[i] + piece_bytes - 1) / piece_bytes; + total_piece_cnt += piece_cnt[i]; + } + for (size_t q = blockIdx.x + static_cast(local_tid) * gridDim.x; q < total_piece_cnt; + q += static_cast(thread_cnt) * gridDim.x) + { + size_t piece_idx = q; + int range_idx = 0; + while (piece_idx >= piece_cnt[range_idx]) + { + piece_idx -= piece_cnt[range_idx]; + range_idx++; + } + size_t const offset = piece_idx * piece_bytes; + size_t const remaining = ranges.bytes[range_idx] - offset; + uint32_t const bytes = static_cast(remaining < piece_bytes ? remaining : piece_bytes) & ~15u; + if (bytes > 0) + { + char const* gmem_ptr = static_cast(ranges.ptr[range_idx]) + offset; + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;\n" ::"l"(gmem_ptr), "r"(bytes) : "memory"); + } + } +} + +// Tensor copies completing on an mbarrier, ldmatrix and HMMA on shared addresses. +__device__ __forceinline__ void mma8_expect_tx(uint64_t* smem_barrier, uint32_t bytes) +{ + asm volatile( + "mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(__nvvm_get_smem_pointer(smem_barrier)), + "r"(bytes) + : "memory"); +} + +__device__ __forceinline__ void mma8_tma_load_2d( + void* smem_dst, CUtensorMap const* tensor_map, int c0, int c1, uint64_t* smem_barrier) +{ + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1, {%2, %3}], [%4];\n" :: + "r"(__nvvm_get_smem_pointer(smem_dst)), + "l"(reinterpret_cast(tensor_map)), "r"(c0), "r"(c1), "r"(__nvvm_get_smem_pointer(smem_barrier)) + : "memory"); +} + +__device__ __forceinline__ void mma8_tma_load_3d( + void* smem_dst, CUtensorMap const* tensor_map, int c0, int c1, int c2, uint64_t* smem_barrier) +{ + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1, {%2, %3, %4}], " + "[%5];\n" ::"r"(__nvvm_get_smem_pointer(smem_dst)), + "l"(reinterpret_cast(tensor_map)), "r"(c0), "r"(c1), "r"(c2), + "r"(__nvvm_get_smem_pointer(smem_barrier)) + : "memory"); +} + +__device__ __forceinline__ void mma8_ldsm_x4(uint32_t smem_addr, uint32_t (&r)[4]) +{ + asm volatile("ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(smem_addr)); +} + +__device__ __forceinline__ void mma8_hmma(float (&d)[4], uint32_t const (&a)[4], uint32_t b0, uint32_t b1) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +} + +__device__ __forceinline__ uint64_t mma8_globaltimer() +{ + uint64_t t; + asm volatile("mov.u64 %0, %%globaltimer;\n" : "=l"(t)); + return t; +} +#endif + +// The GLM-5 fused-A GEMM: gemm_k 6144, 16 weight rows per CTA, up to 8 tokens, with all eight warps issuing the +// HMMAs. output: [gemm_n, gemm_m]; tma_a: the tensor map of the weights [gemm_m, gemm_k]; tma_b: that of the +// activations [gemm_n, gemm_k] (makeFusedAMma8Map). +// +// Stage s is the k range [256 s, 256 s + 256); warp w owns stages w, w + 8 and w + 16 and waits only on its own data. +// Before the grid dependency wait (no kernel of the step writes the weights): the weight tile of stage w is copied into +// shared memory (one 3D tensor copy: four 64-wide k boxes of the CTA's 16 rows, 128-byte swizzle), the weights of +// stages w + 8 and w + 16 reach registers, and the l2_prefetch ranges go to L2. After the wait: the activation tile of +// each of the warp's stages (one 3D copy at 8 tokens, else four 2D boxes of gemm_n rows; tile rows past gemm_n only +// feed HMMA columns that are never stored). The dependent launch is triggered right after the wait. +// +// The register stages permute the k order inside each 32-wide k pair alike for both operands: lane (g = lane / 4, +// c = lane % 4) holds k [8 c, 8 c + 8) of weight rows g and g + 8 and reads k [8 c, 8 c + 8) of token g from the +// activation tile; HMMA step t of the pair takes words 2 t and 2 t + 1 of each. Each warp keeps two accumulators +// (16-wide k step i of a stage adds into i % 2); the eight warp partials are summed in warp order. The products are +// those of a plain GEMM, summed in this order in FP32 and rounded to BF16 once. +// +// The weights of stages w + 8 and w + 16 come in through shared memory, still before the wait: two more 3D tensor +// copies per warp into the space of the activation tiles (those land only after the wait) extended by 32 KB, read into +// registers once they land. Tensor copies leave the SM's load path to a co-resident kernel (the post-MoE all-reduce hop +// that fused-A follows). Lane 0 of warp w issues its two register-stage copies w * kFusedAMma8PaceNs after its stage-w +// copy (%globaltimer), and the l2_prefetch ranges go out only once the CTA's staged tiles are in registers: the CTAs' +// weights reach HBM over 7 * kFusedAMma8PaceNs ahead of the ranges, and the hop beside fused-A keeps more of the HBM +// bandwidth. Issue times do not change the outputs. +template +__global__ __launch_bounds__(256, 1) void fused_a_gemm_mma8_kernel(bf16_t* output, int gemm_n, + FusedAL2Prefetch l2_prefetch_ranges, const __grid_constant__ CUtensorMap tma_a, + const __grid_constant__ CUtensorMap tma_b) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + constexpr int gemm_k = kFusedAHiddenSize; + constexpr int warp_cnt = 8; + constexpr int tile_m = kFusedARowsPerCta; + constexpr int tile_n = kFusedAMaxTokens; + constexpr int tile_k = 256; + constexpr int box_k = 64; + constexpr int boxes = tile_k / box_k; + constexpr int pair_cnt = tile_k / 32; + constexpr int acc_cnt = 2; + static_assert(gemm_k == 3 * warp_cnt * tile_k, "warp w owns stages w, w + 8 and w + 16"); + static_assert(gemm_m % tile_m == 0); + constexpr uint32_t a_box_bytes = tile_m * box_k * sizeof(bf16_t); + constexpr uint32_t b_box_bytes = tile_n * box_k * sizeof(bf16_t); + constexpr uint32_t a_tile_bytes = boxes * a_box_bytes; + constexpr uint32_t b_tile_bytes = boxes * b_box_bytes; + + // 1024-byte aligned (the 128-byte swizzle atoms): stage barriers, 8 weight tiles, 24 activation tiles. Before the + // wait, the 16 register-stage weight tiles sit where the activation tiles land after it (warp w's stages w + 8 and + // w + 16 at tiles 2 w and 2 w + 1), barrier 24 + w for warp w's two. + extern __shared__ char smem[]; + char* smem_base = smem + ((1024u - (__nvvm_get_smem_pointer(smem) & 1023u)) & 1023u); + uint64_t* full = reinterpret_cast(smem_base); + uint64_t* staged = full + 3 * warp_cnt; + char* smem_a = smem_base + 1024; + char* smem_b = smem_a + warp_cnt * a_tile_bytes; + + int const warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + int const lane = threadIdx.x % 32; + int const m0 = blockIdx.x * tile_m; + int const g = lane / 4; + int const c = lane % 4; + uint32_t const b_tx_bytes = boxes * gemm_n * box_k * sizeof(bf16_t); + bool const b_3d = gemm_n == tile_n; + + if (threadIdx.x == 0) + { + for (int i = 0; i < 3 * warp_cnt; i++) + { + // Stage s's barrier: the weight and the activation copies (stages 0-7), the activation copy (8-23). + initialize_barrier(full + i, i < warp_cnt ? 2 : 1); + } + for (int i = 0; i < warp_cnt; i++) + { + initialize_barrier(staged + i, 1); + } + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + } + __syncthreads(); + + if (lane == 0) + { + asm volatile("prefetch.tensormap [%0];\n" ::"l"(reinterpret_cast(&tma_a)) : "memory"); + asm volatile("prefetch.tensormap [%0];\n" ::"l"(reinterpret_cast(&tma_b)) : "memory"); + mma8_expect_tx(full + warp, a_tile_bytes); + mma8_tma_load_3d(smem_a + warp * a_tile_bytes, &tma_a, 0, m0, warp * boxes, full + warp); + mma8_expect_tx(staged + warp, 2 * a_tile_bytes); + uint64_t const due = mma8_globaltimer() + static_cast(warp) * kFusedAMma8PaceNs; + while (mma8_globaltimer() < due) + { + } +#pragma unroll + for (int j = 0; j < 2; j++) + { + mma8_tma_load_3d(smem_b + (2 * warp + j) * a_tile_bytes, &tma_a, 0, m0, (warp + warp_cnt * (j + 1)) * boxes, + staged + warp); + } + } + uint4 a_reg[2][pair_cnt][2]; + // Every weight copy of the CTA is queued ahead of its share of the L2 prefetch list. + __syncthreads(); + // The staged tiles into registers: lane (g, c) takes k [8 c, 8 c + 8) + 32 p of rows g and g + 8, i.e. 16-byte + // chunk (p % 2) * 4 + c of box p / 2 (128-byte rows, 128-byte swizzle). + wait_barrier(staged + warp, 0); +#pragma unroll + for (int j = 0; j < 2; j++) + { + uint32_t const tile = __nvvm_get_smem_pointer(smem_b + (2 * warp + j) * a_tile_bytes); +#pragma unroll + for (int p = 0; p < pair_cnt; p++) + { +#pragma unroll + for (int h = 0; h < 2; h++) + { + uint4& v = a_reg[j][p][h]; + asm volatile("ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "r"(tile + (p / 2) * a_box_bytes + (g + 8 * h) * 128 + ((((p % 2) * 4 + c) ^ g) << 4))); + } + } + } + // The activation copies after the wait overwrite the staged tiles. + __syncthreads(); + if constexpr (l2_prefetch) + { + prefetch_ranges_l2(l2_prefetch_ranges, threadIdx.x, blockDim.x); + } + cudaGridDependencySynchronize(); + cudaTriggerProgrammaticLaunchCompletion(); + + if (lane == 0) + { + // The generic-proxy reads of the staged tiles before the tensor copies that overwrite them. + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); +#pragma unroll + for (int j = 0; j < 3; j++) + { + int const stage = warp + warp_cnt * j; + mma8_expect_tx(full + stage, b_tx_bytes); + if (b_3d) + { + mma8_tma_load_3d(smem_b + stage * b_tile_bytes, &tma_b, 0, 0, stage * boxes, full + stage); + } + else + { +#pragma unroll + for (int q = 0; q < boxes; q++) + { + mma8_tma_load_2d(smem_b + stage * b_tile_bytes + q * b_box_bytes, &tma_b, + stage * tile_k + q * box_k, 0, full + stage); + } + } + } + } + + float acc[acc_cnt][4] = {}; + // Stage w from shared memory. ldmatrix lanes: weight row (lane % 16) at k (lane / 16) * 8 of a 16-wide step; + // token (lane % 8) at k (lane / 8) * 8 of a pair of steps. + { + wait_barrier(full + warp, 0); + int const a_row = lane % 16; + int const a_k = (lane / 16) * 8; + int const b_row = lane % 8; + int const b_k = (lane / 8) * 8; + uint32_t const a_base = __nvvm_get_smem_pointer(smem_a + warp * a_tile_bytes); + uint32_t const b_base = __nvvm_get_smem_pointer(smem_b + warp * b_tile_bytes); +#pragma unroll + for (int s = 0; s < tile_k / 16; s += 2) + { + int const q = s / 4; + int const kk = (s % 4) * 16; + uint32_t b[4]; + mma8_ldsm_x4(b_base + q * b_box_bytes + b_row * 128 + ((((kk + b_k) >> 3) ^ (b_row & 7)) << 4), b); +#pragma unroll + for (int t = 0; t < 2; t++) + { + uint32_t a[4]; + mma8_ldsm_x4( + a_base + q * a_box_bytes + a_row * 128 + ((((kk + 16 * t + a_k) >> 3) ^ (a_row & 7)) << 4), a); + mma8_hmma(acc[(s + t) % acc_cnt], a, b[2 * t], b[2 * t + 1]); + } + } + } + // Stages w + 8 and w + 16: weights from registers, activations from shared memory in the same permuted k order. +#pragma unroll + for (int j = 0; j < 2; j++) + { + int const stage = warp + warp_cnt * (j + 1); + wait_barrier(full + stage, 0); + uint32_t const b_base = __nvvm_get_smem_pointer(smem_b + stage * b_tile_bytes); +#pragma unroll + for (int p = 0; p < pair_cnt; p++) + { + int const chunk = (p % 2) * 4 + c; + uint4 bq; + asm volatile("ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(bq.x), "=r"(bq.y), "=r"(bq.z), "=r"(bq.w) + : "r"(b_base + (p / 2) * b_box_bytes + g * 128 + ((chunk ^ (g & 7)) << 4))); + uint32_t const a0[4] = {a_reg[j][p][0].x, a_reg[j][p][1].x, a_reg[j][p][0].y, a_reg[j][p][1].y}; + uint32_t const a1[4] = {a_reg[j][p][0].z, a_reg[j][p][1].z, a_reg[j][p][0].w, a_reg[j][p][1].w}; + mma8_hmma(acc[(2 * p) % acc_cnt], a0, bq.x, bq.y); + mma8_hmma(acc[(2 * p + 1) % acc_cnt], a1, bq.z, bq.w); + } + } + + // Warp partials into the warp's consumed stage w + 8 activation tile (no copy is pending there), then the sum + // over the warps in warp order. Output (row r, token n) is register (n % 2) + 2 (r / 8) of lane + // (r % 8) * 4 + n / 2. + __syncwarp(); + float* partial = reinterpret_cast(smem_b + (warp + warp_cnt) * b_tile_bytes); +#pragma unroll + for (int r = 0; r < 4; r++) + { + partial[lane * 4 + r] = acc[0][r] + acc[1][r]; + } + __syncthreads(); + if (threadIdx.x < tile_m * tile_n) + { + int const r = threadIdx.x % tile_m; + int const n = threadIdx.x / tile_m; + int const src = ((r % 8) * 4 + n / 2) * 4 + (n % 2) + 2 * (r / 8); + float v = 0.f; +#pragma unroll + for (int w = 0; w < warp_cnt; w++) + { + v += reinterpret_cast(smem_b + (w + warp_cnt) * b_tile_bytes)[src]; + } + if (n < gemm_n) + { + output[static_cast(n) * gemm_m + m0 + r] = __float2bfloat16_rn(v); + } + } +#endif +} + +// Tensor map of a row-major BF16 [rows, cols] matrix with the 128-byte swizzle: dims (k within a 64-wide box, row, box) +// and box (64, box_rows, 4) when box3d (one copy loads [4 boxes][box_rows][64]), else dims (k, row) and box +// (64, box_rows). +CUtensorMap makeFusedAMma8Map(void const* ptr, int rows, int cols, int box_rows, bool box3d) +{ + CUtensorMap map; + cuuint64_t const dims3[3] = {64, static_cast(rows), static_cast(cols / 64)}; + cuuint64_t const strides3[2] = {static_cast(cols) * sizeof(bf16_t), 64 * sizeof(bf16_t)}; + cuuint32_t const box3[3] = {64, static_cast(box_rows), 4}; + cuuint64_t const dims2[2] = {static_cast(cols), static_cast(rows)}; + cuuint64_t const strides2[1] = {static_cast(cols) * sizeof(bf16_t)}; + cuuint32_t const box2[2] = {64, static_cast(box_rows)}; + cuuint32_t const elem_strides[3] = {1, 1, 1}; + TLLM_CU_CHECK(tensorrt_llm::common::CUDADriverWrapper::getInstance()->cuTensorMapEncodeTiled(&map, + CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, box3d ? 3 : 2, const_cast(ptr), box3d ? dims3 : dims2, + box3d ? strides3 : strides2, box3d ? box3 : box2, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, + CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_L2_256B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return map; +} + +// weight: [gemm_m, 6144]; hiddenStates: [numTokens, 6144], numTokens in [1, 8]. +template +void launchFusedAMma8(bf16_t* output, bf16_t const* hiddenStates, bf16_t const* weight, int numTokens, + FusedAL2Prefetch const* l2Prefetch, cudaStream_t stream) +{ + // Alignment slack, stage barriers, 8 weight tiles of 16 x 256, and the space of the 24 activation tiles of 8 x 256 + // extended to the 16 register-stage weight tiles of 16 x 256. + constexpr int kSmemBytes = 1024 + 1024 + 8 * 16 * 256 * sizeof(bf16_t) + 16 * 16 * 256 * sizeof(bf16_t); + CUtensorMap const tmaA = makeFusedAMma8Map(weight, gemm_m, kFusedAHiddenSize, kFusedARowsPerCta, true); + CUtensorMap const tmaB + = makeFusedAMma8Map(hiddenStates, numTokens, kFusedAHiddenSize, numTokens, numTokens == kFusedAMaxTokens); + auto* kernel + = l2Prefetch != nullptr ? &fused_a_gemm_mma8_kernel : &fused_a_gemm_mma8_kernel; + TLLM_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes)); + cudaLaunchConfig_t config; + config.gridDim = dim3(gemm_m / kFusedARowsPerCta, 1, 1); + config.blockDim = dim3(256); + config.dynamicSmemBytes = kSmemBytes; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = attrs; + TLLM_CUDA_CHECK(cudaLaunchKernelEx( + &config, kernel, output, numTokens, l2Prefetch != nullptr ? *l2Prefetch : FusedAL2Prefetch{}, tmaA, tmaB)); +} + +} // namespace + +void invokeFusedAGemm(__nv_bfloat16* output, __nv_bfloat16 const* hiddenStates, __nv_bfloat16 const* weight, + int numTokens, int outFeatures, FusedAL2Prefetch const* l2Prefetch, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(numTokens >= 1 && numTokens <= kFusedAMaxTokens, + "The GLM fused-A GEMM takes 1 to %d tokens, got %d", kFusedAMaxTokens, numTokens); + switch (outFeatures) + { + case kFusedAOutFeatures: + launchFusedAMma8(output, hiddenStates, weight, numTokens, l2Prefetch, stream); + break; + case kFusedAIndexerOutFeatures: + launchFusedAMma8(output, hiddenStates, weight, numTokens, l2Prefetch, stream); + break; + default: + TLLM_THROW("The GLM fused-A GEMM serves %d or %d output features, got %d", kFusedAOutFeatures, + kFusedAIndexerOutFeatures, outFeatures); + } +} + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.h new file mode 100644 index 000000000000..b198da9087ba --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.h @@ -0,0 +1,64 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +//! The GEMM's K: the GLM-5 hidden size, 24 k stages of 256, three per MMA warp. +constexpr int kFusedAHiddenSize = 6144; +//! Rows of the GLM-5 kv_a_proj_with_mqa weight: q_a (2048), compressed KV (512) and k_pe (64). +constexpr int kFusedAOutFeatures = 2624; +//! Rows of the same weight with the DSA indexer's wk (128) and weights_proj (32) rows appended. +constexpr int kFusedAIndexerOutFeatures = 2784; +//! Weight rows per CTA: the grid is outFeatures / 16 CTAs (164 or 174), one per SM. +constexpr int kFusedARowsPerCta = 16; +//! Tokens per call at most: one activation tile holds 8 token rows. +constexpr int kFusedAMaxTokens = 8; + +//! Byte ranges that the GEMM prefetches into L2 before its grid-dependency wait: weights that later kernels of the +//! step read. Unused slots have zero bytes. +struct FusedAL2Prefetch +{ + static constexpr int kMaxRanges = 4; + void const* ptr[kMaxRanges]; + size_t bytes[kMaxRanges]; +}; + +//! output[t, r] = sum_k hiddenStates[t, k] * weight[r, k]: BF16 operands, FP32 accumulation, BF16 output, all three +//! dense row-major. hiddenStates is [numTokens, 6144] with numTokens in [1, 8]; weight is [outFeatures, 6144] with +//! outFeatures kFusedAOutFeatures or kFusedAIndexerOutFeatures; output is [numTokens, outFeatures]. +//! +//! Launched with the programmatic-dependent-launch attribute when TRTLLM_ENABLE_PDL is set. Before its grid-dependency +//! wait the kernel reads only the weight (tensor copies into shared memory and registers) and prefetches the +//! l2Prefetch ranges into L2 (none when l2Prefetch is null); it reads hiddenStates and writes output after the wait, +//! and triggers its dependents right after the wait. Nothing launched before it in the step may write the weight. +void invokeFusedAGemm(__nv_bfloat16* output, __nv_bfloat16 const* hiddenStates, __nv_bfloat16 const* weight, + int numTokens, int outFeatures, FusedAL2Prefetch const* l2Prefetch, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.cu new file mode 100644 index 000000000000..282fd4fe89a7 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.cu @@ -0,0 +1,95 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/envUtils.h" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ +namespace +{ + +constexpr int kThreadsPerBlock = 128; +// BF16 elements per 16-byte vector. +constexpr int kElemsPerVec = 8; + +// blockIdx.x selects the token, blockIdx.y the 128-vector slice of its row; each thread copies one vector. +template +__global__ void __launch_bounds__(kThreadsPerBlock) fusedEmbeddingLookupKernel(uint4* __restrict__ out, + IdT const* __restrict__ ids, uint4 const* __restrict__ table, int64_t vocabSize, int32_t vecsPerRow) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + // The ids may be written by the preceding kernel. Dependents read the output only after their own wait. + cudaGridDependencySynchronize(); + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + int32_t const vec = static_cast(blockIdx.y * kThreadsPerBlock + threadIdx.x); + if (vec >= vecsPerRow) + { + return; + } + int64_t const token = blockIdx.x; + int64_t const id = static_cast(ids[token]); + + // Ids outside [0, vocabSize), such as the -1 padding of rejected draft tokens, give a zero row. + uint4 value = make_uint4(0U, 0U, 0U, 0U); + if (id >= 0 && id < vocabSize) + { + value = table[id * vecsPerRow + vec]; + } + out[token * vecsPerRow + vec] = value; +} + +} // namespace + +void invokeFusedEmbeddingLookup(__nv_bfloat16* out, void const* ids, bool idsAreInt64, __nv_bfloat16 const* table, + int64_t numTokens, int64_t vocabSize, int64_t hiddenSize, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(hiddenSize % kElemsPerVec == 0, "fusedEmbeddingLookup needs a hidden size divisible by 8"); + if (numTokens == 0) + { + return; + } + + int32_t const vecsPerRow = static_cast(hiddenSize / kElemsPerVec); + int32_t const blocksPerRow = (vecsPerRow + kThreadsPerBlock - 1) / kThreadsPerBlock; + dim3 const grid(static_cast(numTokens), static_cast(blocksPerRow)); + dim3 const block(kThreadsPerBlock); + auto* outVec = reinterpret_cast(out); + auto const* tableVec = reinterpret_cast(table); + if (idsAreInt64) + { + tensorrt_llm::common::launchWithPdlWhenEnabled("fusedEmbeddingLookup", fusedEmbeddingLookupKernel, + grid, block, 0, stream, outVec, static_cast(ids), tableVec, vocabSize, vecsPerRow); + } + else + { + tensorrt_llm::common::launchWithPdlWhenEnabled("fusedEmbeddingLookup", fusedEmbeddingLookupKernel, + grid, block, 0, stream, outVec, static_cast(ids), tableVec, vocabSize, vecsPerRow); + } +} + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.h new file mode 100644 index 000000000000..d15db15454f9 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.h @@ -0,0 +1,46 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +//! Embedding lookup of an unsharded BF16 table with the out-of-range guard of the Embedding module: +//! out[t, :] = (0 <= ids[t] < vocabSize) ? table[ids[t], :] : 0 +//! +//! One thread copies one 16-byte vector of a row; the grid covers numTokens x (hiddenSize / 8) vectors. +//! Under PDL the kernel waits for its producer before reading the ids and then lets its dependents launch. +//! +//! \param out Output [numTokens, hiddenSize], row-major, 16-byte aligned. +//! \param ids Token ids [numTokens], int64 when idsAreInt64, else int32. +//! \param table Embedding table [vocabSize, hiddenSize], row-major, 16-byte aligned. +//! \param hiddenSize Row length in elements; must be a multiple of 8. +void invokeFusedEmbeddingLookup(__nv_bfloat16* out, void const* ids, bool idsAreInt64, __nv_bfloat16 const* table, + int64_t numTokens, int64_t vocabSize, int64_t hiddenSize, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLamport.cuh b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLamport.cuh new file mode 100644 index 000000000000..a5147530112e --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLamport.cuh @@ -0,0 +1,182 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/common/lamportUtils.cuh" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +//! \brief The Lamport flags of an MNNVL workspace, read and rotated as common::LamportFlags does. +//! +//! The flags are nine 32-bit words: [0] the current buffer, [1] the dirty buffer (the previous call's), [2] bytes per +//! buffer, [3] stages of the dirty buffer, [4..7] bytes to clear per stage of the dirty buffer, [8] the arrival +//! counter. A kernel publishes into the current buffer, empties the dirty one, and the thread that rotates the flags +//! makes the current buffer the dirty one once every CTA has counted its arrival. Each CTA counts its arrival with a +//! CTA barrier and a release reduction, as common::LamportFlags does, so kernels of either type follow each other on +//! one workspace. On top of common::LamportFlags: a volatile read of the flags, the arrival counter, the rotation +//! without its wait, and the fill of the dirty buffer with a value other than -0.0. +//! +//! Like common::LamportFlags, the clear and the fill need every thread of a 1D-block grid. +template +class GlmLamportFlags +{ +public: + //! \param bufferFlags The nine flag words of the workspace. + //! \param numStages Stages this kernel uses in the current buffer. + //! \param volatileRead Read the flags with volatile loads, which do not go through L1: for a kernel that reads them + //! before its grid dependency wait. + __device__ GlmLamportFlags(uint32_t* bufferFlags, uint32_t numStages, bool volatileRead) + : mBufferFlagsPtr(bufferFlags) + , mFlagAccessPtr(&bufferFlags[kArrivalCounterWord]) + { + uint4 flag; + uint4 clear; + if (volatileRead) + { + auto const flagWords = common::loadPackedVolatile(bufferFlags); + auto const clearWords = common::loadPackedVolatile(bufferFlags + 4); + flag = make_uint4(flagWords.words[0], flagWords.words[1], flagWords.words[2], flagWords.words[3]); + clear = make_uint4(clearWords.words[0], clearWords.words[1], clearWords.words[2], clearWords.words[3]); + } + else + { + flag = reinterpret_cast(bufferFlags)[0]; + clear = reinterpret_cast(bufferFlags)[1]; + } + mCurrentIndex = flag.x; + mDirtyIndex = flag.y; + mCurBufferLayout.numStages = numStages; + mCurBufferLayout.bytesPerBuffer = flag.z; + mDirtyBufferLayout.bytesPerBuffer = flag.z; + mDirtyBufferLayout.numStages = flag.w; + mBytesToClear[0] = clear.x; + mBytesToClear[1] = clear.y; + mBytesToClear[2] = clear.z; + mBytesToClear[3] = clear.w; + } + + //! The base of stage stageIdx of the current buffer. + [[nodiscard]] __device__ void* getCurLamportBuf(void* bufferBasePtr, int stageIdx = 0) const + { + return mCurBufferLayout.getStagePtr(bufferBasePtr, mCurrentIndex, stageIdx); + } + + //! Writes the Lamport empty value (-0.0) into every stage of the dirty buffer. + __device__ void clearDirtyLamportBuf(void* bufferBasePtr) + { + fillDirtyLamportBuf(bufferBasePtr, common::getPackedLamportInit()); + } + + //! Writes fillValue into every stage of the dirty buffer, for a workspace whose empty value is not -0.0. + __device__ void fillDirtyLamportBuf(void* bufferBasePtr, PackedType fillValue) + { + uint32_t const globalCtaIdx = blockIdx.x * gridDim.y + blockIdx.y; + uint32_t const globalTid = globalCtaIdx * blockDim.x + threadIdx.x; + uint32_t const numThreads = gridDim.x * gridDim.y * blockDim.x; + for (uint32_t stage = 0; stage < mDirtyBufferLayout.numStages; stage++) + { + uint32_t const fillBoundary = (mBytesToClear[stage] + sizeof(PackedType) - 1) / sizeof(PackedType); + auto* dirty = static_cast(mDirtyBufferLayout.getStagePtr(bufferBasePtr, mDirtyIndex, stage)); + for (uint32_t packedIdx = globalTid; packedIdx < fillBoundary; packedIdx += numThreads) + { + dirty[packedIdx] = fillValue; + } + } + } + + //! Counts this CTA's arrival: a barrier of the CTA's threads (named barrier 1), then a release reduction on the + //! arrival counter. Every flag read of the CTA precedes it. For kernels of SM 90 and newer. + __device__ void ctaArrive() + { + uint32_t const barrierThreads + = ((static_cast(blockDim.x) + common::kWARP_SIZE - 1U) / common::kWARP_SIZE) * common::kWARP_SIZE; + if (threadIdx.x < common::kWARP_SIZE) + { + asm volatile("barrier.cta.sync 1, %0;" ::"r"(barrierThreads) : "memory"); + if (threadIdx.x == 0) + { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) + asm volatile("red.async.release.global.gpu.add.u32 [%0], %1;" ::"l"(mFlagAccessPtr), "r"(1) : "memory"); +#else + asm volatile("red.release.global.gpu.add.u32 [%0], %1;" ::"l"(mFlagAccessPtr), "r"(1) : "memory"); +#endif + } + } + else + { + asm volatile("barrier.cta.arrive 1, %0;" ::"r"(barrierThreads) : "memory"); + } + } + + //! The arrival counter that ctaArrive increments. + [[nodiscard]] __device__ uint32_t const* arrivalCounter() const + { + return mFlagAccessPtr; + } + + //! The rotation of waitAndUpdate without its wait: the current buffer becomes the dirty one. Only for the thread + //! that waitAndUpdate elects, once it has observed every CTA's arrival with acquire semantics. + __device__ void update(uint4 bytesToClearPerStage) + { + uint4* flagPtr = reinterpret_cast(mBufferFlagsPtr); + flagPtr[0] = {(mCurrentIndex + 1) % common::LamportBufferLayout::sNumLamportBuffers, mCurrentIndex, + mCurBufferLayout.bytesPerBuffer, mCurBufferLayout.numStages}; + flagPtr[1] = bytesToClearPerStage; + *mFlagAccessPtr = 0; + } + + //! Thread 0 of CTA 0 waits until every CTA has arrived, then rotates the flags. + __device__ void waitAndUpdate(uint4 bytesToClearPerStage) + { + if (threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0) + { + uint32_t const targetCount = gridDim.x * gridDim.y * gridDim.z; + uint32_t arrivedCount = 0; + do + { + asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" + : "=r"(arrivedCount) + : "l"(mFlagAccessPtr) + : "memory"); + } while (arrivedCount < targetCount); + update(bytesToClearPerStage); + } + } + +private: + static constexpr int kArrivalCounterWord = 8; + + uint32_t* mBufferFlagsPtr; + uint32_t* mFlagAccessPtr; + uint32_t mCurrentIndex; + uint32_t mDirtyIndex; + uint32_t mBytesToClear[4]; + common::LamportBufferLayout mCurBufferLayout; + common::LamportBufferLayout mDirtyBufferLayout; +}; + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.cu new file mode 100644 index 000000000000..bc67707178a8 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.cu @@ -0,0 +1,838 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/kernels/moeTopKFuncs.cuh" + +#include +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ +namespace +{ +namespace cg = cooperative_groups; + +using Bf16 = __nv_bfloat16; + +constexpr int kNumExperts = kLowMBf16MoeNumExperts; +constexpr int kTopK = kLowMBf16MoeTopK; +static_assert(kNumExperts % 32 == 0, "routing gives each lane of one warp kNumExperts / 32 candidates"); + +// FC1: 8 consumer warps and one producer warp per CTA. Consumer thread t reads the 16-byte pieces t, +// t + kFc1ConsumerThreads, ... of x and of every weight row, so the hidden size is a multiple of kHiddenGranule and at +// most kMaxPiecesPerThread of them. +constexpr int kFc1ConsumerWarps = 8; +constexpr int kFc1ConsumerThreads = kFc1ConsumerWarps * 32; +constexpr int kFc1Threads = kFc1ConsumerThreads + 32; +constexpr int kPieceElems = 8; +constexpr int kHiddenGranule = kFc1ConsumerThreads * kPieceElems; +constexpr int kMaxPiecesPerThread = 4; +// The FC1 ring holds whole weight rows, as many as fit in kFc1RingBytes and at most kMaxStages. At hidden size 6144 +// and interSize 256 that is 7 rows of 12,288 B, and an FC1 CTA (ring, row sums, barriers and the 1,024 B the runtime +// reserves per CTA) takes 87,840 B, so one FC2 CTA (138,312 B) fits beside it in the 233,472 B of an SM. +constexpr int kFc1RingBytes = 7 * 6144 * static_cast(sizeof(Bf16)); +constexpr int kMaxStages = 8; + +// FC2: a warp per routed expert. +constexpr int kFc2Threads = 256; +constexpr int kFc2Warps = kFc2Threads / 32; +static_assert(kTopK == kFc2Warps, "the FC2 kernel gives each routed expert one warp"); + +// FC2: mma.sync m16n8k16 tiles have 16 matrix rows; one 32-element K step is two mma calls (see mmaStep). A warp per +// routed expert walks the intermediate dimension in pairs of K steps, one per accumulator. +constexpr int kMmaRows = 16; +constexpr int kStepK = 32; +constexpr int kFc2StepPairK = 2 * kStepK; +// The largest intermediate size per rank whose FC2 working set fits in shared memory (kTiles = 1). +constexpr int kMaxInterSize = 512; + +// The routed scaling factor as the FC2 kernel takes it: the double value of noaux_tc_op's FP64 weight expression, its +// float value, and whether routingWeight's FP32 path applies to it (decided on the host, see makeRoutingScale). +struct RoutingScale +{ + double value; + float valueFloat; + int fastPath; +}; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +constexpr int kCandidates = kNumExperts / 32; +constexpr int kFc1ProducerWarp = kFc1ConsumerWarps; +// Named barrier at which the FC1 consumer warps wait while the producer selects the experts and issues the first rows. +constexpr int kFc1ParkBarrier = 1; +// The FC2 kernel holds its routing computation until this long after its entry; see lowMBf16MoeFc2Kernel. +constexpr uint64_t kFc2RouteHoldNs = 1000; + +using TopKKey = reduce_topk::TopKRedType; + +// Same expression as sigmoid_accurate in noAuxTcKernels.cu. +__device__ __forceinline__ float sigmoidAccurate(float x) +{ + return 0.5f * tanhf(0.5f * x) + 0.5f; +} + +__device__ __forceinline__ uint32_t warpMax(uint32_t value) +{ + uint32_t result; + asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n" : "=r"(result) : "r"(value)); + return result; +} + +// The correction bias of the experts a lane holds (i * 32 + lane) as float. A constant, so the kernels may read it +// before their grid-dependency wait. +template +__device__ __forceinline__ void loadBias(BiasT const* bias, int lane, float (&values)[kCandidates]) +{ +#pragma unroll + for (int i = 0; i < kCandidates; ++i) + { + values[i] = static_cast(bias[i * 32 + lane]); + } +} + +// The routing keys of noaux_tc_op for n_group == 1 (deepseek_v3_topk_kernel, single group) of the experts a lane +// holds: the sigmoid score plus the correction bias, with the expressions of that kernel, packed with the expert index +// as reduceTopK (moeTopKFuncs.cuh) packs them. Stores the sigmoid scores when sigmoidScores is not null. The logits +// are read through L2. +__device__ __forceinline__ void buildKeys(float const* routerLogits, float const (&biasValues)[kCandidates], int lane, + TopKKey (&keys)[kCandidates], float* sigmoidScores) +{ +#pragma unroll + for (int i = 0; i < kCandidates; ++i) + { + int const expert = i * 32 + lane; + float const sigmoid = sigmoidAccurate(__ldcg(routerLogits + expert)); + if (sigmoidScores != nullptr) + { + sigmoidScores[expert] = sigmoid; + } + keys[i] = TopKKey{sigmoid + biasValues[i], expert}; + } +} + +// The top kTopK of one full warp's keys in the order of reduceTopK: descending score, on equal scores the lower index +// first. The per-lane sort is reduceTopK's; each round then takes the warp maximum of the lanes' leading keys with two +// 32-bit reductions, of the score half and then of the index half among the lanes holding the maximum score, which is +// the maximum of reduceTopK's 64-bit reduction. So the experts and their order are noaux_tc_op's. Returns the k-th +// expert on lane k < kTopK, 0 on the other lanes. +__device__ __forceinline__ int32_t selectTopK(TopKKey (&keys)[kCandidates], int lane) +{ + reduce_topk::Sort::run(keys); + + int32_t laneExpert = 0; + uint64_t packedMax = 0; +#pragma unroll + for (int k = 0; k < kTopK; ++k) + { + if (k > 0 && packedMax == keys[0].compValIdx) + { + // This lane held the previous maximum: drop it, as reduceTopK does. +#pragma unroll + for (int i = 0; i + 1 < kCandidates; ++i) + { + keys[i] = keys[i + 1]; + } + keys[kCandidates - 1] = TopKKey{-INFINITY, (kCandidates - 1) * 32 + lane}; + } + uint32_t const scoreBits = static_cast(keys[0].compValIdx >> 32); + uint32_t const maxScoreBits = warpMax(scoreBits); + uint32_t const indexBits = warpMax(scoreBits == maxScoreBits ? static_cast(keys[0].compValIdx) : 0U); + packedMax = (static_cast(maxScoreBits) << 32) | indexBits; + if (lane == k) + { + laneExpert = TopKKey::kMaxIdx - static_cast(indexBits & 0xFFFF); + } + } + return laneExpert; +} + +// noaux_tc_op's routing weight, static_cast(s * c / (S + 1e-20)) in double with s the expert's sigmoid score, +// S the sum of the selected scores and c the routed scaling factor, without FP64 instructions on the common path: FP64 +// in the FC2 CTAs drops the SM clock for the whole FC1 window (measured on Rubin: 2.38 -> 1.6 GHz). When c is a float +// with at most 20 significant bits (fastPath) and S >= 2^-13, the double expression is the float rounding of the exact +// quotient s * c / S: s * c is exact in double, S + 1e-20 rounds to S, and rounding to double first cannot move a +// quotient of these operands across a float rounding boundary except at an exact midpoint, which rounds the same +// either way. The quotient here comes from an exact two-float product, a correctly rounded division and its exact +// remainder; it is returned when it lies more than 2^-20 ulp inside its rounding interval (and is not a power of two, +// where the interval is lopsided). Otherwise, and for scores below 2^-60 (where the FP32 intermediates could lose +// bits), the FP64 expression runs, out of line so that none of it is scheduled on the common path. +__device__ __noinline__ float routingWeightFp64(float s, float S, double c) +{ + return static_cast(s * c / (S + 1e-20)); +} + +__device__ __forceinline__ float routingWeight(float s, float S, RoutingScale const& scale) +{ + if (scale.fastPath != 0 && S >= 0x1p-13f && s >= 0x1p-60f) + { + float const c = scale.valueFloat; + float const pHi = __fmul_rn(s, c); + float const pLo = __fmaf_rn(s, c, -pHi); + float const q = __fdiv_rn(pHi, S); + float const r = __fmaf_rn(-q, S, pHi); + float const e = __fdiv_rn(__fadd_rn(r, pLo), S); + float const y = __fadd_rn(q, e); + uint32_t const bits = __float_as_uint(y); + if (y > 0.F && (bits & 0x7FFFFFU) != 0 && (bits & 0x7F800000U) != 0x7F800000U) + { + float const ulp = __uint_as_float(bits & 0x7F800000U) * 0x1p-23f; + float const g = __fsub_rn(e, __fsub_rn(y, q)); + if (fabsf(g) < 0.5F * ulp - 0x1p-20f * ulp) + { + return y; + } + } + } + return routingWeightFp64(s, S, scale.value); +} + +__device__ __forceinline__ uint32_t smemAddress(void const* ptr) +{ + return static_cast(__cvta_generic_to_shared(ptr)); +} + +__device__ __forceinline__ void mbarrierInit(uint64_t* bar, uint32_t count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(smemAddress(bar)), "r"(count) : "memory"); +} + +// Makes the barrier initialization visible to the bulk-copy (async) proxy. +__device__ __forceinline__ void fenceBarrierInit() +{ + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); +} + +__device__ __forceinline__ void mbarrierArriveExpectTx(uint64_t* bar, uint32_t bytes) +{ + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(smemAddress(bar)), "r"(bytes) + : "memory"); +} + +// Arrives on the barrier with release semantics: this thread's earlier shared-memory reads are done before the phase +// can complete. +__device__ __forceinline__ void mbarrierArrive(uint64_t* bar) +{ + asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];\n" ::"r"(smemAddress(bar)) : "memory"); +} + +// Whether the barrier's phase with the given parity has completed (one try, which may suspend the thread for a +// while). +__device__ __forceinline__ bool mbarrierTryWait(uint64_t* bar, uint32_t parity) +{ + uint32_t done = 0; + asm volatile( + "{\n" + ".reg .pred complete;\n" + "mbarrier.try_wait.parity.shared::cta.b64 complete, [%1], %2;\n" + "selp.u32 %0, 1, 0, complete;\n" + "}\n" + : "=r"(done) + : "r"(smemAddress(bar)), "r"(parity) + : "memory"); + return done != 0; +} + +__device__ __forceinline__ void mbarrierWait(uint64_t* bar, uint32_t parity) +{ + while (!mbarrierTryWait(bar, parity)) + { + } +} + +// Waits at named barrier `id` until `threads` threads (whole warps) have arrived. +__device__ __forceinline__ void namedBarrierSync(int id, int threads) +{ + asm volatile("bar.sync %0, %1;\n" ::"r"(id), "r"(threads) : "memory"); +} + +__device__ __forceinline__ uint64_t globalTimerNs() +{ + uint64_t now; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(now)); + return now; +} + +// Streamed weights are read once per call: mark their L2 lines evict-first. +__device__ __forceinline__ uint64_t evictFirstPolicy() +{ + uint64_t policy; + asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;\n" : "=l"(policy)); + return policy; +} + +// 1D bulk copy global -> shared that completes `bytes` transactions on `bar`. +__device__ __forceinline__ void bulkCopyToShared( + void* dst, void const* src, uint32_t bytes, uint64_t* bar, uint64_t policy) +{ + uint64_t const srcGlobal = static_cast(__cvta_generic_to_global(src)); + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], %4;\n" + : + : "r"(smemAddress(dst)), "l"(srcGlobal), "r"(bytes), "r"(smemAddress(bar)), "l"(policy) + : "memory"); +} + +__device__ __forceinline__ void mmaBf16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) +{ + asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +// One 32-element K step of a 16-row tile times a vector. Lane (g = lane / 4, t = lane % 4) passes elements [8t, 8t + 8) +// of the step from tile row g (lo), tile row g + 8 (hi) and the vector (vec). The fragment's K columns +// {2t, 2t + 1, 2t + 8, 2t + 9} take the lane's elements {0, 1, 2, 3} in the first call and {4, 5, 6, 7} in the second: +// one K permutation for the matrix and the vector, which leaves the dot products unchanged. Every column of the B +// fragment is the vector, so all accumulator elements hold dot products: elements 0 and 1 of row g, 2 and 3 of row +// g + 8. +__device__ __forceinline__ void mmaStep(float (&acc)[4], uint4 const& lo, uint4 const& hi, uint4 const& vec) +{ + mmaBf16(acc, lo.x, hi.x, lo.y, hi.y, vec.x, vec.y); + mmaBf16(acc, lo.z, hi.z, lo.w, hi.w, vec.z, vec.w); +} + +__device__ __forceinline__ float roundToBf16(float value) +{ + return __bfloat162float(__float2bfloat16_rn(value)); +} + +// The FP32 values of the eight BF16 elements of a 16-byte piece, in memory order. +__device__ __forceinline__ void unpackPiece(uint4 const& piece, float (&values)[kPieceElems]) +{ + uint32_t const words[4] = {piece.x, piece.y, piece.z, piece.w}; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + values[2 * i] = __uint_as_float(words[i] << 16); + values[2 * i + 1] = __uint_as_float(words[i] & 0xFFFF0000U); + } +} + +// Dot product of a 16-byte piece of a BF16 weight row with the FP32 values of the same piece of x, in element order +// with FP32 fused multiply-adds. +__device__ __forceinline__ float dotPiece(uint4 const& piece, float const (&values)[kPieceElems]) +{ + float weights[kPieceElems]; + unpackPiece(piece, weights); + float dot = weights[0] * values[0]; +#pragma unroll + for (int e = 1; e < kPieceElems; ++e) + { + dot = fmaf(weights[e], values[e], dot); + } + return dot; +} + +// The warp sums of two values, each by an xor butterfly over offsets 16, 8, 4, 2, 1, the two shuffle chains +// interleaved. +__device__ __forceinline__ void warpSum2(float& a, float& b) +{ +#pragma unroll + for (int offset = 16; offset > 0; offset /= 2) + { + a += __shfl_xor_sync(0xffffffffU, a, offset); + b += __shfl_xor_sync(0xffffffffU, b, offset); + } +} +#endif + +// FC1 + SwiGLU, one CTA per SM (the grid is the SM count). The work is kTopK * interSize units; unit k * interSize + j +// is intermediate index j of routed slot k, that is the expert's gate row (W1, w3w1 row interSize + j) and up row (W3, +// row j). CTA b takes units [b * U / G, (b + 1) * U / G) and streams their rows, gate then up, each as one bulk copy +// into a ring of numStages row buffers. The producer warp selects the experts and issues the rows; it refills a buffer +// as soon as every consumer warp has read its pieces of it, with no CTA-wide barrier per row. The consumer warps wait +// at a named barrier while the producer selects the experts and issues the first rows. Each consumer thread keeps its +// pieces of x in registers and adds its shares of the dot products of two rows at a time, so that the latency chains of +// the two overlap; each consumer warp reduces its shares with shuffles into shared memory, and each unit's SwiGLU adds +// the 8 warp sums in a fixed order. +template +__global__ void __launch_bounds__(kFc1Threads, 2) + lowMBf16MoeFc1Kernel(Bf16* __restrict__ act, Bf16 const* __restrict__ x, float const* __restrict__ routerLogits, + BiasT const* __restrict__ bias, Bf16 const* __restrict__ w3w1, int hiddenSize, int interSize, int numStages) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + extern __shared__ __align__(128) uint8_t smem[]; + // rowFull[s]: buffer s holds its next row (the producer's arrival plus the copy's bytes). rowFree[s]: every + // consumer warp has read its pieces of buffer s. + __shared__ uint64_t rowFull[kMaxStages]; + __shared__ uint64_t rowFree[kMaxStages]; + __shared__ int32_t topExperts[kTopK]; + + int const warp = threadIdx.x / 32; + int const lane = threadIdx.x % 32; + int const rowBytes = hiddenSize * static_cast(sizeof(Bf16)); + int const numUnits = kTopK * interSize; + int const unitBegin = static_cast(static_cast(blockIdx.x) * numUnits / gridDim.x); + int const unitEnd = static_cast(static_cast(blockIdx.x + 1) * numUnits / gridDim.x); + int const numRows = 2 * (unitEnd - unitBegin); + uint8_t* ring = smem; + float* rowSums = reinterpret_cast(smem + numStages * rowBytes); + + if (threadIdx.x == 0) + { + for (int s = 0; s < numStages; ++s) + { + mbarrierInit(&rowFull[s], 1); + mbarrierInit(&rowFree[s], kFc1ConsumerWarps); + } + fenceBarrierInit(); + } + // The correction bias is a constant: the producer reads it before the grid-dependency wait. + float biasValues[kCandidates] = {}; + if (warp == kFc1ProducerWarp) + { + loadBias(bias, lane, biasValues); + } + __syncthreads(); + + // x and the router logits are written by earlier kernels. + cudaGridDependencySynchronize(); + // The FC2 kernel reads only the logits, the bias and its weights before its own wait, and the logits are complete + // since the wait above returned. Released here, its CTAs take the room beside this kernel's and stage their + // weights while this kernel streams. + cudaTriggerProgrammaticLaunchCompletion(); + + if (warp == kFc1ProducerWarp) + { + TopKKey keys[kCandidates]; + buildKeys(routerLogits, biasValues, lane, keys, nullptr); + int32_t const expert = selectTopK(keys, lane); + if (lane < kTopK) + { + topExperts[lane] = expert; + } + __syncwarp(); + + // Lane 0 streams the rows. The next row is of unit (slot, index), the gate row first, into buffer s, which it + // fills for the pass-th time. + uint64_t const policy = evictFirstPolicy(); + int64_t const expertStride = 2 * static_cast(interSize) * hiddenSize; + int slot = unitBegin / interSize; + int index = unitBegin % interSize; + bool gateRow = true; + Bf16 const* expertRows = w3w1 + topExperts[slot] * expertStride; + int s = 0; + uint32_t pass = 0; + auto issueRow = [&]() + { + int64_t const rowInExpert = gateRow ? interSize + index : index; + mbarrierArriveExpectTx(&rowFull[s], rowBytes); + bulkCopyToShared(ring + s * rowBytes, expertRows + rowInExpert * hiddenSize, rowBytes, &rowFull[s], policy); + if (!gateRow && ++index == interSize) + { + // The up row closed the slot's last unit. + index = 0; + if (++slot < kTopK) + { + expertRows = w3w1 + topExperts[slot] * expertStride; + } + } + gateRow = !gateRow; + if (++s == numStages) + { + s = 0; + ++pass; + } + }; + int const firstRows = numStages < numRows ? numStages : numRows; + if (lane == 0) + { + for (int row = 0; row < firstRows; ++row) + { + issueRow(); + } + } + // Release the consumer warps. + __syncwarp(); + namedBarrierSync(kFc1ParkBarrier, kFc1Threads); + if (lane == 0) + { + for (int row = firstRows; row < numRows; ++row) + { + // Buffer s still holds row - numStages until every consumer warp has read it. + mbarrierWait(&rowFree[s], (pass - 1) & 1U); + issueRow(); + } + } + } + else + { + float xValues[kMaxPiecesPerThread][kPieceElems] = {}; + int const piecesPerThread = hiddenSize / kHiddenGranule; +#pragma unroll + for (int i = 0; i < kMaxPiecesPerThread; ++i) + { + if (i < piecesPerThread) + { + unpackPiece( + __ldcg(reinterpret_cast(x) + threadIdx.x + i * kFc1ConsumerThreads), xValues[i]); + } + } + // Wait at the barrier, without polling the ring, while the producer selects the experts and issues the first + // rows. + namedBarrierSync(kFc1ParkBarrier, kFc1Threads); + int s0 = 0; + uint32_t pass0 = 0; + // Rows row and row + 1 (numRows is even) are in buffers s0 and s1 (numStages >= 2), filled for the pass0-th and + // pass1-th time. + for (int row = 0; row < numRows; row += 2) + { + int s1 = s0 + 1; + uint32_t pass1 = pass0; + if (s1 == numStages) + { + s1 = 0; + ++pass1; + } + mbarrierWait(&rowFull[s0], pass0 & 1U); + mbarrierWait(&rowFull[s1], pass1 & 1U); + uint4 const* pieces0 = reinterpret_cast(ring + s0 * rowBytes); + uint4 const* pieces1 = reinterpret_cast(ring + s1 * rowBytes); + float dot0 = dotPiece(pieces0[threadIdx.x], xValues[0]); + float dot1 = dotPiece(pieces1[threadIdx.x], xValues[0]); +#pragma unroll + for (int i = 1; i < kMaxPiecesPerThread; ++i) + { + if (i < piecesPerThread) + { + dot0 += dotPiece(pieces0[threadIdx.x + i * kFc1ConsumerThreads], xValues[i]); + dot1 += dotPiece(pieces1[threadIdx.x + i * kFc1ConsumerThreads], xValues[i]); + } + } + // This warp is done with both buffers (the arrival releases its lanes' reads, ordered before it by the + // warp barrier); the producer refills a buffer once every consumer warp is. + __syncwarp(); + if (lane == 0) + { + mbarrierArrive(&rowFree[s0]); + mbarrierArrive(&rowFree[s1]); + } + warpSum2(dot0, dot1); + if (lane == 0) + { + rowSums[row * kFc1ConsumerWarps + warp] = dot0; + rowSums[(row + 1) * kFc1ConsumerWarps + warp] = dot1; + } + s0 = s1 + 1; + pass0 = pass1; + if (s0 == numStages) + { + s0 = 0; + ++pass0; + } + } + } + __syncthreads(); + + for (int unit = static_cast(threadIdx.x); unit < numRows / 2; unit += kFc1Threads) + { + float gate = 0.F; + float up = 0.F; + for (int w = 0; w < kFc1ConsumerWarps; ++w) + { + gate += rowSums[2 * unit * kFc1ConsumerWarps + w]; + up += rowSums[(2 * unit + 1) * kFc1ConsumerWarps + w]; + } + // The CUTLASS FC1 GEMM writes its output in BF16 and the activation kernel reads it back. + gate = roundToBf16(gate); + up = roundToBf16(up); + // silu(gate) * up in FP32 as the CUTLASS activation (GLUAdaptor) computes it, then BF16. + float const silu = gate * (1.F / (1.F + expf(-gate))); + act[unitBegin + unit] = __float2bfloat16_rn(silu * up); + } +#endif +} + +// FC2 + the weighted sum over the routed experts. Grid: one CTA per kTiles * 16 output rows; warp k computes routed +// expert k's terms from its weight slice, staged in shared memory before the grid-dependency wait. Warp 0 routes before +// that wait, from the router logits and the bias staged by bulk copies, without FP64 instructions. After the wait every +// warp stages its expert's SwiGLU outputs itself. +template +__global__ void __launch_bounds__(kFc2Threads, 2) + lowMBf16MoeFc2Kernel(Bf16* __restrict__ out, Bf16 const* __restrict__ act, float const* __restrict__ routerLogits, + BiasT const* __restrict__ bias, Bf16 const* __restrict__ w2, int hiddenSize, int interSize, RoutingScale scale) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + constexpr int kRows = kTiles * kMmaRows; + extern __shared__ __align__(128) uint8_t smem[]; + __shared__ __align__(16) float logitsShared[kNumExperts]; + __shared__ __align__(16) BiasT biasShared[kNumExperts]; + __shared__ uint64_t inputsFull; + __shared__ uint64_t slicesFull; + __shared__ float sigmoidScores[kNumExperts]; + __shared__ int32_t topExperts[kTopK]; + __shared__ float topWeights[kTopK]; + __shared__ float terms[kTopK][kRows]; + + uint64_t const entryNs = globalTimerNs(); + + int const sliceElems = kRows * interSize; + Bf16* slices = reinterpret_cast(smem); + Bf16* actShared = slices + kTopK * sliceElems; + + int const warp = threadIdx.x / 32; + int const lane = threadIdx.x % 32; + int const row0 = blockIdx.x * kRows; + + if (threadIdx.x == 0) + { + mbarrierInit(&inputsFull, 1); + mbarrierInit(&slicesFull, 1); + fenceBarrierInit(); + } + __syncthreads(); + + if (warp == 0) + { + // Before the grid-dependency wait this kernel reads only the router logits, the bias and the FC2 weights. The + // logits come from the router GEMM, which completed before the FC1 kernel returned from its own wait and + // released this grid; the bias and the weights are constants. The logits and the bias come in by bulk copies: + // plain loads issued now queue behind the FC1 kernel's bulk copies on this SM for microseconds (measured). + bool const bulkInputs + = ((reinterpret_cast(routerLogits) | reinterpret_cast(bias)) % 16) == 0; + if (bulkInputs) + { + if (lane == 0) + { + uint64_t const policy = evictFirstPolicy(); + mbarrierArriveExpectTx(&inputsFull, sizeof(logitsShared) + sizeof(biasShared)); + bulkCopyToShared(logitsShared, routerLogits, sizeof(logitsShared), &inputsFull, policy); + bulkCopyToShared(biasShared, bias, sizeof(biasShared), &inputsFull, policy); + } + mbarrierWait(&inputsFull, 0); + } + else + { + // Each lane stages, and below reads back, only the experts it holds. +#pragma unroll + for (int i = 0; i < kCandidates; ++i) + { + logitsShared[i * 32 + lane] = __ldcg(routerLogits + i * 32 + lane); + biasShared[i * 32 + lane] = bias[i * 32 + lane]; + } + } + // The FC1 kernel's producer warps select the same experts on these SMs at about this time, on the FC1 + // kernel's critical path, and routing here concurrently makes their selection about twice as long (measured + // 0.5 -> 1.1 us). So the routing waits until kFc2RouteHoldNs after this CTA's entry; the weight copies are + // still issued long before the FC1 kernel ends. + while (globalTimerNs() < entryNs + kFc2RouteHoldNs) + { + __nanosleep(64); + } + // The FC1 kernel's routing (the keys of buildKeys, selectTopK) on the staged inputs, and the weights from + // routingWeight: the experts, their order and the FP32 weights of noaux_tc_op. + TopKKey keys[kCandidates]; +#pragma unroll + for (int i = 0; i < kCandidates; ++i) + { + int const expert = i * 32 + lane; + float const sigmoid = sigmoidAccurate(logitsShared[expert]); + sigmoidScores[expert] = sigmoid; + keys[i] = TopKKey{sigmoid + static_cast(biasShared[expert]), expert}; + } + int32_t const expert = selectTopK(keys, lane); + __syncwarp(); + auto const tile = cg::tiled_partition<32>(cg::this_thread_block()); + float const scoreNorm = lane < kTopK ? sigmoidScores[expert] : 0.F; + float const sumNorm = cg::reduce(tile, scoreNorm, cg::plus{}); + if (lane < kTopK) + { + topExperts[lane] = expert; + topWeights[lane] = routingWeight(scoreNorm, sumNorm, scale); + } + __syncwarp(); + if (lane == 0) + { + mbarrierArriveExpectTx(&slicesFull, kTopK * sliceElems * static_cast(sizeof(Bf16))); + } + __syncwarp(); + if (lane < kTopK) + { + bulkCopyToShared(slices + lane * sliceElems, + w2 + (static_cast(topExperts[lane]) * hiddenSize + row0) * interSize, + sliceElems * static_cast(sizeof(Bf16)), &slicesFull, evictFirstPolicy()); + } + } + __syncthreads(); + + // The SwiGLU outputs are written by the FC1 kernel: every warp waits for it, then stages its expert's outputs. + cudaGridDependencySynchronize(); + Bf16 const* slice = slices + warp * sliceElems; + Bf16* expertAct = actShared + warp * interSize; + for (int i = lane; i < interSize / 8; i += 32) + { + reinterpret_cast(expertAct)[i] = __ldcg(reinterpret_cast(act + warp * interSize) + i); + } + __syncwarp(); + mbarrierWait(&slicesFull, 0); + + // Lane (g, t) of warp k accumulates the dot products of slice rows tile * 16 + g and tile * 16 + g + 8 with + // act[k, :], over the two 32-element halves of each 64-element K step in two accumulators. Lanes of odd g take the + // two halves in the other order, so that a warp's loads of 8 rows spread over all 32 banks: column n of the B + // fragments then holds the vector in the order of the rows of n's parity, and the accumulators of those columns + // hold the same products in the same order as with one order for all (which of the two holds which half differs). + int const g = lane / 4; + int const t = lane % 4; + int const swap = g & 1; + float acc[kTiles][2][4] = {}; + for (int j0 = 0; j0 < interSize; j0 += kFc2StepPairK) + { +#pragma unroll + for (int half = 0; half < 2; ++half) + { + int const j = j0 + (half ^ swap) * kStepK + 8 * t; + uint4 const av = *reinterpret_cast(expertAct + j); +#pragma unroll + for (int tileIdx = 0; tileIdx < kTiles; ++tileIdx) + { + uint4 const lo = *reinterpret_cast(slice + (tileIdx * kMmaRows + g) * interSize + j); + uint4 const hi = *reinterpret_cast(slice + (tileIdx * kMmaRows + g + 8) * interSize + j); + mmaStep(acc[tileIdx][half], lo, hi, av); + } + } + } + + // The CUTLASS FC2 epilogue rounds each expert's scaled output to BF16 before adding it into the output. Lane (g, 0) + // holds columns 0 and 1 of rows g and g + 8, of which column g & 1 has their dot products. + float const weight = topWeights[warp]; + if (t == 0) + { +#pragma unroll + for (int tileIdx = 0; tileIdx < kTiles; ++tileIdx) + { + float const lo0 = swap ? acc[tileIdx][0][1] : acc[tileIdx][0][0]; + float const lo1 = swap ? acc[tileIdx][1][1] : acc[tileIdx][1][0]; + float const hi0 = swap ? acc[tileIdx][0][3] : acc[tileIdx][0][2]; + float const hi1 = swap ? acc[tileIdx][1][3] : acc[tileIdx][1][2]; + terms[warp][tileIdx * kMmaRows + g] = roundToBf16(weight * (lo0 + lo1)); + terms[warp][tileIdx * kMmaRows + g + 8] = roundToBf16(weight * (hi0 + hi1)); + } + } + __syncthreads(); + if (threadIdx.x < kRows) + { + float sum = 0.F; +#pragma unroll + for (int k = 0; k < kTopK; ++k) + { + sum += terms[k][threadIdx.x]; + } + out[row0 + threadIdx.x] = __float2bfloat16_rn(sum); + } +#endif +} + +// Lets a kernel use smemBytes of dynamic shared memory and asks for the largest shared-memory carveout for both +// kernels, so that an SM running a CTA of one can take a CTA of the other. +template +void setSharedMemoryAttributes(Kernel* kernel, size_t smemBytes) +{ + TLLM_CUDA_CHECK( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smemBytes))); + TLLM_CUDA_CHECK(cudaFuncSetAttribute( + kernel, cudaFuncAttributePreferredSharedMemoryCarveout, static_cast(cudaSharedmemCarveoutMaxShared))); +} + +// routingWeight's FP32 path applies when the scaling factor is a float with at most 20 significant bits (the low 4 +// mantissa bits zero) between 2^-20 and 2^20. +RoutingScale makeRoutingScale(double routedScalingFactor) +{ + float const valueFloat = static_cast(routedScalingFactor); + uint32_t bits = 0; + std::memcpy(&bits, &valueFloat, sizeof(bits)); + bool const fastPath = static_cast(valueFloat) == routedScalingFactor && (bits & 0xFU) == 0 + && valueFloat >= 0x1p-20f && valueFloat <= 0x1p20f; + return RoutingScale{routedScalingFactor, valueFloat, fastPath ? 1 : 0}; +} + +template +void launchFc2(Bf16* out, Bf16 const* act, float const* routerLogits, BiasT const* bias, Bf16 const* w2, int hiddenSize, + int interSize, double routedScalingFactor, cudaStream_t stream) +{ + int const rows = kTiles * kMmaRows; + size_t const smemBytes = static_cast(kTopK) * (rows + 1) * interSize * sizeof(Bf16); + auto* kernel = &lowMBf16MoeFc2Kernel; + setSharedMemoryAttributes(kernel, smemBytes); + common::launchWithPdlWhenEnabled("lowMBf16MoeFc2", kernel, dim3(hiddenSize / rows), dim3(kFc2Threads), smemBytes, + stream, out, act, routerLogits, bias, w2, hiddenSize, interSize, makeRoutingScale(routedScalingFactor)); +} + +} // namespace + +bool lowMBf16MoeSupportsShape(int64_t hiddenSize, int64_t interSize) +{ + return hiddenSize > 0 && hiddenSize % kHiddenGranule == 0 && hiddenSize <= kMaxPiecesPerThread * kHiddenGranule + && interSize > 0 && interSize % kFc2StepPairK == 0 && interSize <= kMaxInterSize; +} + +template +void invokeLowMBf16Moe(Bf16* out, Bf16* act, Bf16 const* x, float const* routerLogits, BiasT const* bias, + Bf16 const* w3w1, Bf16 const* w2, int hiddenSize, int interSize, double routedScalingFactor, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(lowMBf16MoeSupportsShape(hiddenSize, interSize), + "lowMBf16Moe: unsupported hidden size %d or intermediate size %d", hiddenSize, interSize); + + // FC1: one CTA per SM, each with an even share of the units and a slot per row and consumer warp for the sums. + static int const smCount = common::getMultiProcessorCount(); + int const numUnits = kTopK * interSize; + int const fc1Grid = std::min(smCount, numUnits); + int const rowBytes = hiddenSize * static_cast(sizeof(Bf16)); + int const numStages = std::min(kMaxStages, kFc1RingBytes / rowBytes); + int const maxRows = 2 * ((numUnits + fc1Grid - 1) / fc1Grid); + size_t const fc1SmemBytes + = static_cast(numStages) * rowBytes + static_cast(maxRows) * kFc1ConsumerWarps * sizeof(float); + auto* fc1 = &lowMBf16MoeFc1Kernel; + setSharedMemoryAttributes(fc1, fc1SmemBytes); + common::launchWithPdlWhenEnabled("lowMBf16MoeFc1", fc1, dim3(fc1Grid), dim3(kFc1Threads), fc1SmemBytes, stream, act, + x, routerLogits, bias, w3w1, hiddenSize, interSize, numStages); + + // Two 16-row tiles per FC2 CTA where the 8 weight slices fit in shared memory, else one. + if (interSize <= kMaxInterSize / 2) + { + launchFc2(out, act, routerLogits, bias, w2, hiddenSize, interSize, routedScalingFactor, stream); + } + else + { + launchFc2(out, act, routerLogits, bias, w2, hiddenSize, interSize, routedScalingFactor, stream); + } +} + +#define INSTANTIATE_LOW_M_BF16_MOE(BiasT) \ + template void invokeLowMBf16Moe(Bf16 * out, Bf16 * act, Bf16 const* x, float const* routerLogits, \ + BiasT const* bias, Bf16 const* w3w1, Bf16 const* w2, int hiddenSize, int interSize, \ + double routedScalingFactor, cudaStream_t stream); + +INSTANTIATE_LOW_M_BF16_MOE(float); +INSTANTIATE_LOW_M_BF16_MOE(__nv_bfloat16); + +#undef INSTANTIATE_LOW_M_BF16_MOE + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.h new file mode 100644 index 000000000000..928eebda65e2 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.h @@ -0,0 +1,74 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +//! Routed experts and experts per token the low-M BF16 MoE kernels are built for. +constexpr int kLowMBf16MoeNumExperts = 256; +constexpr int kLowMBf16MoeTopK = 8; + +//! Whether invokeLowMBf16Moe supports a layer: hiddenSize a multiple of 2048 and at most 8192, interSize a multiple +//! of 64 and at most 512. +bool lowMBf16MoeSupportsShape(int64_t hiddenSize, int64_t interSize); + +//! One token through an MoE layer with unquantized BF16 SwiGLU experts held as tensor-parallel slices (no expert +//! parallelism) and DeepSeek-V3 routing with one expert group, in two kernels: +//! 1. route, then for each routed expert e_k: act[k, :] = bf16(silu(bf16(W1[e_k] x)) * bf16(W3[e_k] x)); +//! 2. route again, then out = bf16(sum_k bf16(w_k * (W2[e_k] act[k, :]))). +//! Routing is the noaux_tc_op computation for n_group == 1 (deepseek_v3_topk_kernel): the same experts, in the same +//! order, and bit-identical weights w_k (normalized sigmoid scores times routedScalingFactor). The BF16 rounding points +//! are those of the CUTLASS MoE path: the FC1 output, the SwiGLU output and each expert's scaled FC2 output. The sum +//! over the experts is in FP32 in a fixed order, where the CUTLASS finalize adds BF16 terms with atomics. +//! +//! The FC1 kernel runs one CTA per SM: a producer warp selects the experts and streams whole weight rows through a +//! shared-memory ring, which 8 consumer warps read two rows at a time. It releases the FC2 kernel right after its +//! grid-dependency wait, so that FC2 CTAs fit beside it and stage their weights in shared memory while FC1 streams. +//! Neither kernel runs FP64 instructions on its common path (FP64 lowers the SM clock): the FC2 kernel forms the +//! weights in FP32 with the result of noaux_tc_op's FP64 expression, which it runs only where that is not certain. +//! +//! Under PDL the FC1 kernel reads x and the logits after its grid-dependency wait; the FC2 kernel reads only the +//! logits, the bias and the FC2 weights before its own wait, which is safe because the FC1 kernel releases it after +//! its wait has returned. +//! +//! \param out Output [hiddenSize]. +//! \param act Scratch [kLowMBf16MoeTopK, interSize] for the SwiGLU outputs, 16-byte aligned. +//! \param x Input token [hiddenSize], 16-byte aligned. +//! \param routerLogits Router logits [kLowMBf16MoeNumExperts], FP32. +//! \param bias Routing correction bias [kLowMBf16MoeNumExperts]. +//! \param w3w1 FC1 weights [kLowMBf16MoeNumExperts, 2 * interSize, hiddenSize], 16-byte aligned: rows +//! [0, interSize) of an expert are W3 (up), rows [interSize, 2 * interSize) are W1 (gate), the +//! layout CutlassFusedMoE stores. +//! \param w2 FC2 weights [kLowMBf16MoeNumExperts, hiddenSize, interSize], 16-byte aligned. +template +void invokeLowMBf16Moe(__nv_bfloat16* out, __nv_bfloat16* act, __nv_bfloat16 const* x, float const* routerLogits, + BiasT const* bias, __nv_bfloat16 const* w3w1, __nv_bfloat16 const* w2, int hiddenSize, int interSize, + double routedScalingFactor, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.cu new file mode 100644 index 000000000000..91246637cf1f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.cu @@ -0,0 +1,231 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ +namespace +{ + +constexpr int kNopeDim = kGlmKvHalfLatentDim; +constexpr int kRopeDim = kGlmKvHalfRopeDim; +// BF16 values per 16-byte vector. +constexpr int kEltsPerVec = 8; +constexpr int kNopeVecs = kNopeDim / kEltsPerVec; +constexpr int kRopeVecs = kRopeDim / kEltsPerVec; +// The cache row in 8-element units. +constexpr int kRowVecs = kNopeVecs + kRopeVecs; +constexpr int kWarpSize = 32; +constexpr int kNopeVecsPerLane = kNopeVecs / kWarpSize; +// Four token rows per block, one warp per row. +constexpr int kBlockSize = 128; +static_assert(kNopeVecs % kWarpSize == 0 && kRopeVecs <= kWarpSize, "the row splits over one warp"); + +inline __device__ float2 bf16x2ToFloat2(uint32_t word) +{ + return make_float2(__uint_as_float(word << 16), __uint_as_float(word & 0xFFFF0000u)); +} + +// Eight FP32 values to saturated E4M3. +inline __device__ uint2 toFp8x8(float const (&v)[kEltsPerVec]) +{ + uint2 out; + auto* pairs = reinterpret_cast<__nv_fp8x2_storage_t*>(&out); +#pragma unroll + for (int j = 0; j < kEltsPerVec / 2; ++j) + { + pairs[j] = __nv_cvt_float2_to_fp8x2(make_float2(v[2 * j], v[2 * j + 1]), __NV_SATFINITE, __NV_E4M3); + } + return out; +} + +// One warp per token row. Lane l holds latent vectors l and l + 32 and, for l < 8, K-RoPE vector l. +// kTriggerAtEntry: the PDL dependents may launch at kernel entry instead of after the stores; they read this kernel's +// outputs only after their own grid-dependency wait. +template +__global__ void __launch_bounds__(kBlockThreads) + glmMlaKvHalfAppendKernel(__nv_bfloat16 const* __restrict__ latent, int latentRowStride, KVBlockArray kvCache, + float2 const* __restrict__ cosSinCache, __nv_bfloat16 const* __restrict__ kvNormWeight, float kvNormEps, + int numTokens, int seqLenQ, int32_t const* __restrict__ kvLens, float* __restrict__ bmm1Scale, + float* __restrict__ bmm2Scale, float softmaxScale) +{ + constexpr int kRowsPerBlock = kBlockThreads / kWarpSize; + int const lane = threadIdx.x % kWarpSize; + int const warpId = threadIdx.x / kWarpSize; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if constexpr (kTriggerAtEntry) + { + cudaTriggerProgrammaticLaunchCompletion(); + } + // The latent row is the previous kernel's output: wait before any global read. + cudaGridDependencySynchronize(); +#endif + + if (blockIdx.x == 0 && threadIdx.x == 0) + { + // The scales of the generation attention over this cache (unit KV scales). + constexpr float kLog2e = 1.4426950408889634074f; + bmm1Scale[0] = softmaxScale; + bmm1Scale[1] = softmaxScale * kLog2e; + bmm2Scale[0] = 1.f; + } + + // The norm weight is row-invariant: two 8-element slices per lane, kept in FP32. + float weight[kNopeVecsPerLane][kEltsPerVec]; +#pragma unroll + for (int i = 0; i < kNopeVecsPerLane; ++i) + { + uint4 const w = *reinterpret_cast(kvNormWeight + (lane + i * kWarpSize) * kEltsPerVec); + uint32_t const words[4] = {w.x, w.y, w.z, w.w}; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float2 const f = bf16x2ToFloat2(words[j]); + weight[i][2 * j] = f.x; + weight[i][2 * j + 1] = f.y; + } + } + + for (int token = warpId + blockIdx.x * kRowsPerBlock; token < numTokens; token += kRowsPerBlock * gridDim.x) + { + int const seq = token / seqLenQ; + int const tokenInSeq = token - seq * seqLenQ; + // Cache slot and RoPE position, as the generation RoPE kernel derives them. + int const tokenKvIdx = kvLens[seq] - seqLenQ + tokenInSeq; + if (tokenKvIdx < 0) + { + continue; + } + __nv_bfloat16 const* row = latent + static_cast(token) * static_cast(latentRowStride); + // Page-table read first: it does not depend on the reduction. + auto* kDst = reinterpret_cast(kvCache.getKBlockPtr(seq, tokenKvIdx)); + + uint4 nope[kNopeVecsPerLane]; +#pragma unroll + for (int i = 0; i < kNopeVecsPerLane; ++i) + { + nope[i] = *reinterpret_cast(row + (lane + i * kWarpSize) * kEltsPerVec); + } + uint4 rope = make_uint4(0, 0, 0, 0); + if (lane < kRopeVecs) + { + rope = *reinterpret_cast(row + kNopeDim + lane * kEltsPerVec); + } + + float x[kNopeVecsPerLane][kEltsPerVec]; + float sumSquares = 0.f; +#pragma unroll + for (int i = 0; i < kNopeVecsPerLane; ++i) + { + uint32_t const words[4] = {nope[i].x, nope[i].y, nope[i].z, nope[i].w}; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float2 const f = bf16x2ToFloat2(words[j]); + x[i][2 * j] = f.x; + x[i][2 * j + 1] = f.y; + sumSquares += f.x * f.x + f.y * f.y; + } + } +#pragma unroll + for (int mask = kWarpSize / 2; mask > 0; mask >>= 1) + { + sumSquares += __shfl_xor_sync(0xFFFFFFFFu, sumSquares, mask); + } + float const rstd = rsqrtf(sumSquares / static_cast(kNopeDim) + kvNormEps); + +#pragma unroll + for (int i = 0; i < kNopeVecsPerLane; ++i) + { + float scaled[kEltsPerVec]; +#pragma unroll + for (int j = 0; j < kEltsPerVec; ++j) + { + scaled[j] = x[i][j] * (rstd * weight[i][j]); + } + int const vec = lane + i * kWarpSize; + *reinterpret_cast(kDst + kvCache.getKVLocalIdx(tokenKvIdx, 0, kRowVecs, vec) * kEltsPerVec) + = toFp8x8(scaled); + } + + if (lane < kRopeVecs) + { + // Pair p of the RoPE slice uses entry kRopeDim * position + p of the table. + float2 const* coef = cosSinCache + static_cast(kRopeDim) * tokenKvIdx + lane * (kEltsPerVec / 2); + uint32_t const words[4] = {rope.x, rope.y, rope.z, rope.w}; + float rotated[kEltsPerVec]; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float2 const v = bf16x2ToFloat2(words[j]); + float2 const cs = coef[j]; + rotated[2 * j] = cs.x * v.x - cs.y * v.y; + rotated[2 * j + 1] = cs.x * v.y + cs.y * v.x; + } + *reinterpret_cast( + kDst + kvCache.getKVLocalIdx(tokenKvIdx, 0, kRowVecs, kNopeVecs + lane) * kEltsPerVec) + = toFp8x8(rotated); + } + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if constexpr (!kTriggerAtEntry) + { + cudaTriggerProgrammaticLaunchCompletion(); + } +#endif +} + +} // namespace + +void invokeGlmMlaKvHalfAppend(__nv_bfloat16 const* latent, int32_t latentRowStride, KVBlockArray const& kvCache, + float2 const* cosSinCache, __nv_bfloat16 const* kvNormWeight, float kvNormEps, int32_t numTokens, int32_t seqLenQ, + int32_t const* kvLens, float* bmm1Scale, float* bmm2Scale, float softmaxScale, bool triggerAtEntry, + cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(latent != nullptr && cosSinCache != nullptr && kvNormWeight != nullptr && kvLens != nullptr + && bmm1Scale != nullptr && bmm2Scale != nullptr, + "The GLM MLA KV append needs the latent rows, the rope table, the norm weight, the KV lengths and the scales"); + TLLM_CHECK_WITH_INFO(latentRowStride >= kNopeDim + kRopeDim && latentRowStride % kEltsPerVec == 0, + "latent row stride %d must cover a %d-column row and be a multiple of %d", latentRowStride, kNopeDim + kRopeDim, + kEltsPerVec); + TLLM_CHECK_WITH_INFO(numTokens >= 0 && seqLenQ > 0 && numTokens % seqLenQ == 0, + "MLA can only support input sequences with the same sequence length."); + + constexpr int kRowsPerBlock = kBlockSize / kWarpSize; + dim3 const grid(static_cast(tensorrt_llm::common::divUp(std::max(numTokens, 1), kRowsPerBlock))); + auto const kernel + = triggerAtEntry ? glmMlaKvHalfAppendKernel : glmMlaKvHalfAppendKernel; + tensorrt_llm::common::launchWithPdlWhenEnabled("glmMlaKvHalfAppend", kernel, grid, kBlockSize, 0, stream, latent, + static_cast(latentRowStride), kvCache, cosSinCache, kvNormWeight, kvNormEps, static_cast(numTokens), + static_cast(seqLenQ), kvLens, bmm1Scale, bmm2Scale, softmaxScale); +} + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.h new file mode 100644 index 000000000000..219dd77ad8f0 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.h @@ -0,0 +1,56 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/kernels/kvCacheUtils.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +// The MLA latent row of GLM-5.2: kv_lora_rank 512 latent columns followed by qk_rope_head_dim 64 K-RoPE columns. +constexpr int32_t kGlmKvHalfLatentDim = 512; +constexpr int32_t kGlmKvHalfRopeDim = 64; + +// The KV half of MLA generation for the latent layout above, in one kernel: RMSNorm of the 512 latent columns of the +// raw kv_a_proj output, GPT-J RoPE of the 64 K-RoPE columns, E4M3 quantization at unit scale straight from FP32, and +// the append to the paged FP8 cache through kvCache, at the slot and RoPE position of the generation RoPE kernel +// (position kvLens[s] - seqLenQ + t for token t of sequence s; a negative position is skipped). Block 0 also writes +// the generation attention's scales: bmm1Scale = {softmaxScale, softmaxScale * log2(e)}, bmm2Scale = {1}. +// - latent: [numTokens] rows of at least 576 BF16 values, latentRowStride elements apart (a multiple of 8, rows +// 16-byte aligned); +// - cosSinCache: per position, 64 (cos, sin) float pairs (the duplicated-data table of the MLA rope), of which pair +// p < 32 rotates K-RoPE columns 2p and 2p + 1; +// - kvNormWeight: the 512 kv_a_layernorm weights. +// PDL: the kernel waits on its grid dependency before any global read. With triggerAtEntry its dependents may launch +// at its entry, otherwise after its stores; either way they read its outputs only after their own grid-dependency +// wait. +void invokeGlmMlaKvHalfAppend(__nv_bfloat16 const* latent, int32_t latentRowStride, KVBlockArray const& kvCache, + float2 const* cosSinCache, __nv_bfloat16 const* kvNormWeight, float kvNormEps, int32_t numTokens, int32_t seqLenQ, + int32_t const* kvLens, float* bmm1Scale, float* bmm2Scale, float softmaxScale, bool triggerAtEntry, + cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.cu new file mode 100644 index 000000000000..c2753b020af7 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.cu @@ -0,0 +1,203 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/common/lamportUtils.cuh" +#include "tensorrt_llm/kernels/glmMoeDsa/glmLamport.cuh" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ +namespace +{ + +constexpr int kAllgatherBlockSize = 512; +// Gathered packets each thread polls together. +constexpr int kAllgatherPacketsInFlight = 2; +// A payload word equal to the empty word is sent as this word: still a pair of BF16 NaNs. +constexpr uint32_t kAllgatherSubstituteWord = kMnnvlExchangeEmptyWord - 1U; +constexpr int kPacketBytes = static_cast(sizeof(float4)); + +union AllgatherPacket +{ + float4 packed; + uint32_t words[4]; +}; + +struct AllgatherKernelParams +{ + float4 const* input; // [numRows, shardPackets] packets of BF16 + float* output; // [numRows, nRanks * shardPackets] packets of BF16, written as FP32 + void* bufferLocal; + void* bufferMcast; + uint32_t* bufferFlags; + int numRows; + int shardPackets; + int nRanks; + int rank; +}; + +__device__ __forceinline__ bool isEmptyPacket(AllgatherPacket const& packet) +{ + bool empty = false; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + empty |= packet.words[i] == kMnnvlExchangeEmptyWord; + } + return empty; +} + +// Writes gathered packet packetIdx (eight BF16 values) to the output as FP32. A BF16 value is the upper half of the +// FP32 value it widens to: shift the bits, NaN payloads included. +__device__ __forceinline__ void storeWidened(float* output, int packetIdx, AllgatherPacket const& packet) +{ + float4* out = reinterpret_cast(output) + 2 * packetIdx; + out[0] = make_float4(__uint_as_float(packet.words[0] << 16), __uint_as_float(packet.words[0] & 0xFFFF0000U), + __uint_as_float(packet.words[1] << 16), __uint_as_float(packet.words[1] & 0xFFFF0000U)); + out[1] = make_float4(__uint_as_float(packet.words[2] << 16), __uint_as_float(packet.words[2] & 0xFFFF0000U), + __uint_as_float(packet.words[3] << 16), __uint_as_float(packet.words[3] & 0xFFFF0000U)); +} + +__global__ void __launch_bounds__(kAllgatherBlockSize) mnnvlAllgatherKernel(AllgatherKernelParams params) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + // The shard is written by the preceding kernel. Dependents read the output only after their own wait. + cudaGridDependencySynchronize(); + cudaTriggerProgrammaticLaunchCompletion(); + + GlmLamportFlags flag(params.bufferFlags, /*numStages=*/1, /*volatileRead=*/false); + auto* stageMcast = static_cast(flag.getCurLamportBuf(params.bufferMcast)); + auto const* stageLocal = static_cast(flag.getCurLamportBuf(params.bufferLocal)); + + int const tid = static_cast(blockIdx.x * blockDim.x + threadIdx.x); + int const numThreads = static_cast(gridDim.x * blockDim.x); + int const rowPackets = params.nRanks * params.shardPackets; + int const shardTotal = params.numRows * params.shardPackets; + int const gatherTotal = params.numRows * rowPackets; + + // Store the shard once through the multicast address, at [row][rank][packet] of the gathered layout. + for (int idx = tid; idx < shardTotal; idx += numThreads) + { + int const row = idx / params.shardPackets; + int const col = idx - row * params.shardPackets; + AllgatherPacket packet; + packet.packed = params.input[idx]; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + if (packet.words[i] == kMnnvlExchangeEmptyWord) + { + packet.words[i] = kAllgatherSubstituteWord; + } + } + stageMcast[row * rowPackets + params.rank * params.shardPackets + col] = packet.packed; + } + + flag.ctaArrive(); + // Empty the buffer of the previous call; it is written again two calls later. + AllgatherPacket empty; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + empty.words[i] = kMnnvlExchangeEmptyWord; + } + flag.fillDirtyLamportBuf(params.bufferLocal, empty.packed); + + // Copy the gathered rows out of the local buffer as the packets of every rank land. + for (int base = tid; base < gatherTotal; base += numThreads * kAllgatherPacketsInFlight) + { + AllgatherPacket packets[kAllgatherPacketsInFlight]; + bool pending[kAllgatherPacketsInFlight]; +#pragma unroll + for (int j = 0; j < kAllgatherPacketsInFlight; ++j) + { + int const idx = base + j * numThreads; + pending[j] = idx < gatherTotal; + if (pending[j]) + { + packets[j].packed = common::loadPackedVolatile(&stageLocal[idx]).packed; + } + } + bool anyPending = true; + while (anyPending) + { + anyPending = false; +#pragma unroll + for (int j = 0; j < kAllgatherPacketsInFlight; ++j) + { + if (!pending[j]) + { + continue; + } + int const idx = base + j * numThreads; + if (isEmptyPacket(packets[j])) + { + packets[j].packed = common::loadPackedVolatile(&stageLocal[idx]).packed; + anyPending = true; + } + else + { + storeWidened(params.output, idx, packets[j]); + pending[j] = false; + } + } + } + } + + flag.waitAndUpdate({static_cast(gatherTotal * kPacketBytes), 0, 0, 0}); +#endif +} + +} // namespace + +void mnnvlAllgather(MnnvlAllgatherParams const& params) +{ + TLLM_CHECK_WITH_INFO(common::getSMVersion() >= 90, "[GLM MNNVL all-gather] requires SM 90 or newer."); + constexpr int kEltsPerPacket = kPacketBytes / static_cast(sizeof(__nv_bfloat16)); + TLLM_CHECK_WITH_INFO(params.shardDim % kEltsPerPacket == 0, + "[GLM MNNVL all-gather] a shard row must be a multiple of %d BF16 elements, got %d", kEltsPerPacket, + params.shardDim); + if (params.numRows == 0 || params.shardDim == 0) + { + return; + } + int const shardPackets = params.shardDim / kEltsPerPacket; + int const gatherTotal = params.numRows * params.nRanks * shardPackets; + static int const kSmCount = common::getMultiProcessorCount(); + int constexpr kPacketsPerBlock = kAllgatherBlockSize * kAllgatherPacketsInFlight; + // At most one CTA per SM: every CTA is resident and stores its part of the shard, so the packets the CTAs of other + // ranks poll for are always on their way. + int const numBlocks = std::max(1, std::min((gatherTotal + kPacketsPerBlock - 1) / kPacketsPerBlock, kSmCount)); + AllgatherKernelParams const kernelParams{static_cast(params.input), params.output, + params.bufferPtrLocal, params.multicastPtr, params.bufferFlags, params.numRows, shardPackets, params.nRanks, + params.rank}; + common::launchWithPdlWhenEnabled("glmMnnvlAllgather", mnnvlAllgatherKernel, dim3(numBlocks), + dim3(kAllgatherBlockSize), 0, params.stream, kernelParams); +} + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.h new file mode 100644 index 000000000000..5abbe8bbe048 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.h @@ -0,0 +1,62 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +//! Value of every 32-bit word of an unwritten slot of an exchange workspace; a new workspace is filled with it. +constexpr uint32_t kMnnvlExchangeEmptyWord = 0xFFFFFFFFU; + +//! \brief The one-shot MNNVL all-gather of BF16 rows along the last dimension, widened to FP32, on an exchange +//! workspace: three Lamport buffers whose unwritten words hold kMnnvlExchangeEmptyWord, and their flags (the layout +//! of GlmLamportFlags). +struct MnnvlAllgatherParams +{ + int nRanks; //!< Ranks of the group + int rank; //!< This rank + int numRows; //!< Rows of the shard + int shardDim; //!< BF16 elements per row of one rank's shard, a multiple of 8 + void const* input; //!< This rank's shard [numRows, shardDim], BF16, contiguous, 16-byte aligned + float* output; //!< The gathered rows [numRows, nRanks * shardDim], FP32, contiguous, 16-byte aligned + void* bufferPtrLocal; //!< This rank's unicast address of the workspace's three Lamport buffers + void* multicastPtr; //!< Multicast address of the same buffers + uint32_t* bufferFlags; //!< Lamport flags of the workspace + cudaStream_t stream; +}; + +//! \brief Gathers every rank's shard into [row][rank][column] order, in one kernel. +//! +//! Every rank stores its shard once through the multicast address at [row][rank][column] of the gathered layout and +//! copies the gathered rows out of its local copy of the buffer as the packets land, widening BF16 to FP32 exactly; +//! it then fills the previous call's buffer with the empty word. The output equals the rank-ordered concatenation of +//! the shards bit for bit, except that a payload word equal to kMnnvlExchangeEmptyWord (a pair of BF16 NaNs with every +//! bit set, which no float conversion produces) arrives as kMnnvlExchangeEmptyWord - 1, still a pair of NaNs. The +//! kernel waits for its producer before it touches any memory, and its PDL dependents launch right after that wait. +void mnnvlAllgather(MnnvlAllgatherParams const& params); + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.cu new file mode 100644 index 000000000000..637fa3c8e6f4 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.cu @@ -0,0 +1,799 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaTypeUtils.cuh" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/common/lamportUtils.cuh" +#include "tensorrt_llm/common/reduceKernelUtils.cuh" +#include "tensorrt_llm/kernels/glmMoeDsa/glmLamport.cuh" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ +namespace +{ + +using common::cuda_cast; + +// Eight BF16 elements, one 16-byte access. +constexpr int kEltsPerThread = static_cast(sizeof(float4) / sizeof(__nv_bfloat16)); +constexpr int kWarpSize = 32; +constexpr int kLog2WarpSize = 5; +constexpr int kLaneIdMask = kWarpSize - 1; +// The fast residual + RMSNorm epilogue takes CTAs of up to four whole warps in clusters of 2 to 8 CTAs. +constexpr int kMaxThreads = 128; +constexpr int kMaxWarps = kMaxThreads / kWarpSize; +constexpr int kMaxClusterSize = 8; +// Ranks summed from zero before their sum joins the total. +constexpr int kRankChunk = 8; +// Lamport data per rank up to which an all-reduce runs one-shot: 1 MiB (the MNNVL sweep's threshold), and on 16 ranks +// 1.5 MiB, which keeps an 8-row hidden-6144 BF16 step one-shot. +constexpr int64_t kOneShotBytes = 64 * 1024 * 8 * 2; +constexpr int64_t kOneShotBytes16Ranks = 3 * 512 * 1024; + +union PackedBf16x8 +{ + float4 packed; + __nv_bfloat16 elements[kEltsPerThread]; +}; + +struct HopKernelParams +{ + __nv_bfloat16* normOut; + __nv_bfloat16* residualOut; + __nv_bfloat16 const* shard; + __nv_bfloat16 const* residualIn; + __nv_bfloat16 const* gamma; + __nv_bfloat16** bufferPtrs; + __nv_bfloat16* mcastPtr; + uint32_t* bufferFlags; + int numTokens; + int tokenDim; + int rank; + float epsilon; +}; + +struct MoeFinalizeKernelParams +{ + __nv_bfloat16 const* fc2Output; + __nv_bfloat16 const* expertWeights; + int32_t const* expandedIdxToPermutedIdx; + __nv_bfloat16 const* addend; + int hiddenDimPadded; +}; + +// Shared memory of the fast epilogue: each thread's residual and gamma vectors, copied before the Lamport exchange; the +// warp sums of squares of every CTA of the cluster, kMaxWarps per CTA; and the transaction barrier that completes once +// all of them have landed. +struct FastEpilogueStorage +{ + float4 residual[kMaxThreads]; + float4 gamma[kMaxThreads]; + float4 warpSums[kMaxClusterSize]; + uint64_t warpSumsBarrier; +}; + +// Only the kernels that call this function allocate the storage. +__device__ __forceinline__ FastEpilogueStorage& fastEpilogueStorage() +{ + __shared__ FastEpilogueStorage storage; + return storage; +} + +__device__ __forceinline__ void copyF4(float4* dstShared, __nv_bfloat16 const* src) +{ + __pipeline_memcpy_async(dstShared, reinterpret_cast(src), sizeof(float4)); +} + +// -0.0 is the empty value of the Lamport buffers: a payload element of -0.0 is sent as +0.0. +__device__ __forceinline__ void sanitizeLamportPayload(PackedBf16x8& value) +{ +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + if (common::isNegZero(value.elements[i])) + { + value.elements[i] = cuda_cast<__nv_bfloat16, float>(0.F); + } + } +} + +__device__ __forceinline__ void accumulatePacked(float (&accum)[kEltsPerThread], PackedBf16x8 const& value) +{ +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + accum[i] += cuda_cast(value.elements[i]); + } +} + +__device__ __forceinline__ uint32_t loadArrivalsRelaxed(uint32_t const* arrivalCounter) +{ + uint32_t arrivals; + asm volatile("ld.relaxed.gpu.global.u32 %0, [%1];" : "=r"(arrivals) : "l"(arrivalCounter) : "memory"); + return arrivals; +} + +// Strong loads at GPU scope (L2, never a line that L1 kept from an earlier use), for reads before the grid dependency +// wait. +__device__ __forceinline__ int32_t loadRelaxedGpu(int32_t const* ptr) +{ + int32_t value; + asm volatile("ld.relaxed.gpu.global.b32 %0, [%1];" : "=r"(value) : "l"(ptr)); + return value; +} + +// The bits of a BF16 value. +__device__ __forceinline__ uint16_t loadRelaxedGpu(__nv_bfloat16 const* ptr) +{ + uint16_t bits; + asm volatile("ld.relaxed.gpu.global.b16 %0, [%1];" : "=h"(bits) : "l"(ptr)); + return bits; +} + +// The one-shot exchange in one poll loop: every iteration loads the slots of all WorldSize ranks together until none +// holds the Lamport empty value. With LocalValue the slot of localRank is not waited for: localValue, the bits this +// kernel stored there, stands in for it. With TrackArrivals a thread with a non-null arrivalCounter reads the counter +// (relaxed) at the start of every iteration, so that the last value it keeps in *arrivals is about as recent as the +// peers' data. The sum: the ranks of each chunk of kRankChunk summed in rank order from zero in FP32, the chunk sums +// added to zero in order, rounded to BF16. The WorldSize values in flight need the registers of a launch bound of +// kMaxThreads. +template +__device__ __forceinline__ PackedBf16x8 reduceOneshotPollAll(PackedBf16x8 const& localValue, int localRank, + __nv_bfloat16 const* stagePtrLocal, int token, int tokenDim, int packedIdx, uint32_t const* arrivalCounter, + uint32_t* arrivals) +{ + constexpr int kChunk = WorldSize < kRankChunk ? WorldSize : kRankChunk; + static_assert(WorldSize % kChunk == 0, "The ranks must split into whole chunks"); + __nv_bfloat16 const* slots = &stagePtrLocal[token * tokenDim * WorldSize + packedIdx * kEltsPerThread]; + PackedBf16x8 values[WorldSize]; + while (true) + { + if constexpr (TrackArrivals) + { + if (arrivalCounter != nullptr) + { + *arrivals = loadArrivalsRelaxed(arrivalCounter); + } + } + bool valid = true; +#pragma unroll + for (int r = 0; r < WorldSize; r++) + { + auto const loaded = common::loadPackedVolatile(&slots[r * tokenDim]); + values[r].packed = loaded.packed; + valid &= (LocalValue && r == localRank) | !common::isLamportDirty(loaded); + } + if (valid) + { + break; + } + } + float accum[kEltsPerThread]; +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + accum[i] = 0.F; + } +#pragma unroll + for (int rankBase = 0; rankBase < WorldSize; rankBase += kChunk) + { + float chunkAccum[kEltsPerThread]; +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + chunkAccum[i] = 0.F; + } +#pragma unroll + for (int rr = 0; rr < kChunk; rr++) + { + int const r = rankBase + rr; + accumulatePacked(chunkAccum, LocalValue && r == localRank ? localValue : values[r]); + } +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + accum[i] += chunkAccum[i]; + } + } + PackedBf16x8 packedAccum; +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + packedAccum.elements[i] = cuda_cast<__nv_bfloat16, float>(accum[i]); + } + return packedAccum; +} + +// Called by one thread per CTA before every thread's relaxed arrive on the cluster barrier: the barrier expects one +// arrival (this one) and the bytes of the warp sums of every CTA of the cluster. That cluster barrier phase publishes +// the initialization; every thread waits for the phase before it sends. +__device__ __forceinline__ void initFastEpilogueBarrier() +{ + namespace cg = cooperative_groups; + uint32_t const bytes + = static_cast(cg::this_cluster().num_blocks() * (blockDim.x >> kLog2WarpSize) * sizeof(float)); + uint32_t const barrier = static_cast(__cvta_generic_to_shared(&fastEpilogueStorage().warpSumsBarrier)); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(barrier) : "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(barrier), "r"(bytes) : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); +} + +// Sum of threadSum over all threads of the cluster. Each warp reduces its sums with the warpReduceSum butterfly (every +// lane ends with the same bits); every warp sends its sum to every CTA of the cluster with st.async, and each CTA waits +// on its transaction barrier until all warp sums have landed. Per CTA the warp sums are added as (w0 + w2) + (w1 + w3) +// without the absent terms, and the CTA sums in block-rank order from zero: the order of blockReduceSum +// followed by the cluster loop of the one-shot all-reduce's default epilogue. Needs whole warps, at most kMaxWarps +// warps per CTA, 2 to kMaxClusterSize CTAs per cluster, the barrier initialized by initFastEpilogueBarrier and the +// cluster barrier phase that published it completed. +__device__ __forceinline__ float clusterSumOfSquaresFast(float threadSum) +{ + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + auto& storage = fastEpilogueStorage(); + int const numBlocks = static_cast(cluster.num_blocks()); + int const blockRank = static_cast(cluster.block_rank()); + int const numWarps = static_cast(blockDim.x >> kLog2WarpSize); + int const warpId = static_cast(threadIdx.x >> kLog2WarpSize); + int const laneId = static_cast(threadIdx.x & kLaneIdMask); + + float const warpSum = common::warpReduceSum(threadSum); + uint32_t const barrier = static_cast(__cvta_generic_to_shared(&storage.warpSumsBarrier)); + if (laneId < numBlocks) + { + // Lane l sends to CTA l. + uint32_t const slot = static_cast( + __cvta_generic_to_shared(reinterpret_cast(storage.warpSums) + blockRank * kMaxWarps + warpId)); + uint32_t peerSlot; + uint32_t peerBarrier; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(peerSlot) : "r"(slot), "r"(laneId)); + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(peerBarrier) : "r"(barrier), "r"(laneId)); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(peerSlot), + "r"(__float_as_uint(warpSum)), "r"(peerBarrier) + : "memory"); + } + uint32_t complete = 0; + while (!complete) + { + asm volatile( + "{\n" + ".reg .pred done;\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 done, [%1], 0;\n" + "selp.u32 %0, 1, 0, done;\n" + "}\n" + : "=r"(complete) + : "r"(barrier) + : "memory"); + } + + float fullSum = 0.F; +#pragma unroll + for (int block = 0; block < kMaxClusterSize; ++block) + { + if (block < numBlocks) + { + float4 const sums = storage.warpSums[block]; + float const even = numWarps > 2 ? sums.x + sums.z : sums.x; + float const odd = numWarps > 3 ? sums.y + sums.w : sums.y; + fullSum += numWarps > 1 ? even + odd : even; + } + } + return fullSum; +} + +// The rotation of the Lamport flags by the thread that GlmLamportFlags::waitAndUpdate elects. arrivals is the last +// value that thread read from the arrival counter with a relaxed load. Once it covers every CTA, the fence turns that +// read into an acquire of every CTA's release arrival and the flags rotate at once; otherwise the thread waits on the +// counter as waitAndUpdate does. +__device__ __forceinline__ void rotateLamportFlags( + GlmLamportFlags& flag, uint32_t arrivals, uint4 bytesToClearPerStage) +{ + uint32_t const numCtas = gridDim.x * gridDim.y * gridDim.z; + if (arrivals >= numCtas) + { + asm volatile("fence.acquire.gpu;" ::: "memory"); + } + else + { + do + { + asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" + : "=r"(arrivals) + : "l"(flag.arrivalCounter()) + : "memory"); + } while (arrivals < numCtas); + } + flag.update(bytesToClearPerStage); +} + +// After the poll of the rotating thread: rotate now if the last count it read covers every CTA; otherwise read the +// counter once more, while the epilogue runs, and leave the rotation to the end of the kernel. +__device__ __forceinline__ void rotateIfAllArrived( + GlmLamportFlags& flag, bool& rotationPending, uint32_t& arrivals, uint4 bytesToClearPerStage) +{ + if (rotationPending) + { + if (arrivals >= gridDim.x * gridDim.y * gridDim.z) + { + rotateLamportFlags(flag, arrivals, bytesToClearPerStage); + rotationPending = false; + } + else + { + arrivals = loadArrivalsRelaxed(flag.arrivalCounter()); + } + } +} + +// The fast residual + RMSNorm epilogue of both hops, after the copies of the residual and gamma have completed: the +// residual added in BF16, the BF16 squares summed in FP32 in element order and over the cluster, (x * rsqrt) * gamma +// rounded to BF16, the residual output stored after the exchange of the sums (the complete-tx of each st.async +// releases every earlier store of the thread), and the rotation of the flags if it is still pending. +__device__ __forceinline__ void hopEpilogue(HopKernelParams const& params, GlmLamportFlags& flag, + PackedBf16x8 packedAccum, PackedBf16x8 const& residualIn, int threadOffset, bool rotationPending, uint32_t arrivals, + uint4 bytesToClearPerStage) +{ +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + packedAccum.elements[i] += residualIn.elements[i]; + } + PackedBf16x8 gamma; + gamma.packed = fastEpilogueStorage().gamma[threadIdx.x]; + float threadSum = 0.F; +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + threadSum += cuda_cast(packedAccum.elements[i] * packedAccum.elements[i]); + } + // Every thread arrived on the cluster barrier before its grid dependency wait; waiting for that phase publishes + // every CTA's transaction barrier before any thread sends. + __cluster_barrier_wait(); + float const fullSum = clusterSumOfSquaresFast(threadSum); + reinterpret_cast(¶ms.residualOut[threadOffset])[0] = packedAccum.packed; + float const rcpRms = rsqrtf(fullSum / params.tokenDim + params.epsilon); +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + packedAccum.elements[i] + = cuda_cast<__nv_bfloat16, float>(cuda_cast(packedAccum.elements[i]) * rcpRms + * cuda_cast(gamma.elements[i])); + } + reinterpret_cast(¶ms.normOut[threadOffset])[0] = packedAccum.packed; + if (rotationPending) + { + rotateLamportFlags(flag, arrivals, bytesToClearPerStage); + } +} + +// One row per cluster, eight elements per thread, the clusters' CTAs covering the row exactly. Each CTA counts its +// own arrival (GlmLamportFlags) and the thread that rotates the flags tracks the arrival counter in its poll. The +// transaction barrier is published with a relaxed cluster arrive before the grid dependency wait. The PDL dependents +// launch right after the grid dependency wait (at the start with a published shard) and must read the outputs only +// after their own wait. +// ShardPublished: the producer of the input has stored this rank's rows into the current Lamport buffer of every rank, +// sanitized, where this kernel would store them, so the kernel neither loads nor stores them and polls every rank's +// slot, its own included. It does not wait for its primary before the exchange: the producer releases it only after its +// own grid dependency wait, so every grid before the producer has completed, the previous hop (the flags, the residual) +// included. The flags and the residual are read with volatile loads, gamma (a weight) with the epilogue's copy. A CTA +// counts its arrival only once all of its slots are in: the rotation after the last arrival must not precede a producer +// CTA's read of the flags. The grid waits for its primary at its end, so that its completion still implies the +// producer's. +template +__global__ void __launch_bounds__(kMaxThreads) mnnvlHopKernel(HopKernelParams params) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + int const packedIdx = static_cast(cluster.thread_rank()); + int const token = static_cast(blockIdx.x); + int const threadOffset = token * params.tokenDim + packedIdx * kEltsPerThread; + + // The transaction barrier is this CTA's shared memory, which the previous grid does not touch: initialize it and + // publish it to the cluster while the grid dependency wait is pending. + if (threadIdx.x == 0) + { + initFastEpilogueBarrier(); + } + asm volatile("barrier.cluster.arrive.relaxed;" ::: "memory"); + if constexpr (!ShardPublished) + { + cudaGridDependencySynchronize(); + } + cudaTriggerProgrammaticLaunchCompletion(); + + GlmLamportFlags flag(params.bufferFlags, /*numStages=*/1, /*volatileRead=*/ShardPublished); + auto* stagePtrMcast = static_cast<__nv_bfloat16*>(flag.getCurLamportBuf(params.mcastPtr)); + auto* stagePtrLocal = static_cast<__nv_bfloat16*>(flag.getCurLamportBuf(params.bufferPtrs[params.rank])); + auto& storage = fastEpilogueStorage(); + + // ==================== Broadcast this rank's rows ==================== + PackedBf16x8 val{}; + [[maybe_unused]] PackedBf16x8 residualPublished; + if constexpr (ShardPublished) + { + residualPublished.packed = common::loadPackedVolatile(¶ms.residualIn[threadOffset]).packed; + copyF4(&storage.gamma[threadIdx.x], ¶ms.gamma[packedIdx * kEltsPerThread]); + } + else + { + val.packed = *reinterpret_cast(¶ms.shard[threadOffset]); + // The residual and gamma do not depend on the exchange: copy them now, so that their latency overlaps the + // broadcast and the poll instead of following the reduction. + copyF4(&storage.residual[threadIdx.x], ¶ms.residualIn[threadOffset]); + copyF4(&storage.gamma[threadIdx.x], ¶ms.gamma[packedIdx * kEltsPerThread]); + sanitizeLamportPayload(val); + reinterpret_cast( + &stagePtrMcast[token * params.tokenDim * WorldSize + params.rank * params.tokenDim])[packedIdx] + = val.packed; + } + __pipeline_commit(); + if constexpr (!ShardPublished) + { + flag.ctaArrive(); + } + // ==================== Empty the previous call's buffer ==================== + flag.clearDirtyLamportBuf(params.bufferPtrs[params.rank]); + + // ==================== Reduction ==================== + // The thread that waitAndUpdate would elect tracks the arrival counter and rotates the flags. + bool rotationPending = threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0; + uint32_t arrivals = 0; + PackedBf16x8 packedAccum; + if constexpr (ShardPublished) + { + packedAccum = reduceOneshotPollAll( + val, params.rank, stagePtrLocal, token, params.tokenDim, packedIdx, nullptr, nullptr); + flag.ctaArrive(); + } + else + { + packedAccum = reduceOneshotPollAll(val, params.rank, stagePtrLocal, token, + params.tokenDim, packedIdx, rotationPending ? flag.arrivalCounter() : nullptr, &arrivals); + } + uint4 const bytesToClear{ + static_cast(params.numTokens * params.tokenDim * WorldSize * sizeof(__nv_bfloat16)), 0, 0, 0}; + rotateIfAllArrived(flag, rotationPending, arrivals, bytesToClear); + + // ==================== Residual + RMSNorm ==================== + // Each thread reads only the copies it issued, so their completion is the only wait needed. + __pipeline_wait_prior(0); + PackedBf16x8 residualIn; + if constexpr (ShardPublished) + { + residualIn = residualPublished; + } + else + { + residualIn.packed = storage.residual[threadIdx.x]; + } + hopEpilogue(params, flag, packedAccum, residualIn, threadOffset, rotationPending, arrivals, bytesToClear); + if constexpr (ShardPublished) + { + // The completion of this grid implies its primary's, as when it waits at its start. + cudaGridDependencySynchronize(); + } +#endif +} + +// The routed experts' finalize step fused into the hop (BF16, top-8): each thread computes its eight elements of the +// finalize output as the TRTLLM-Gen finalize does (the token's eight FC2 rows weighted in k order with an FMA, -1 slots +// skipped, rounded to BF16; the addend added in FP32 and the sum rounded once more) and publishes them as its rows. +// Everything else is mnnvlHopKernel without a published shard. Eight threads per CTA read the token's FC2 row indices +// and routing weights before the grid dependency wait: the MoE's routing wrote them, and it has completed when this +// grid launches because FC1 and FC2 trigger their dependents only after their own waits. The launch bound lets the +// eight FC2 rows of a thread be in flight in registers. +template +__global__ void __launch_bounds__(kMaxThreads) + mnnvlMoeFinalizeHopKernel(HopKernelParams params, MoeFinalizeKernelParams moe) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + int const packedIdx = static_cast(cluster.thread_rank()); + int const token = static_cast(blockIdx.x); + int const threadOffset = token * params.tokenDim + packedIdx * kEltsPerThread; + __shared__ int32_t sPermutedIdx[kMnnvlHopMoeTopK]; + __shared__ uint16_t sExpertWeightBits[kMnnvlHopMoeTopK]; + + // The transaction barrier is this CTA's shared memory: initialize it and publish it to the cluster while the grid + // dependency wait is pending. + if (threadIdx.x == 0) + { + initFastEpilogueBarrier(); + } + asm volatile("barrier.cluster.arrive.relaxed;" ::: "memory"); + bool const isLoader = threadIdx.x < kMnnvlHopMoeTopK; + int32_t permutedIdxPart = -1; + uint16_t expertWeightBitsPart = 0; + if (isLoader) + { + int const expandedIdx = token * kMnnvlHopMoeTopK + static_cast(threadIdx.x); + permutedIdxPart = loadRelaxedGpu(&moe.expandedIdxToPermutedIdx[expandedIdx]); + expertWeightBitsPart = loadRelaxedGpu(&moe.expertWeights[expandedIdx]); + } + cudaGridDependencySynchronize(); + cudaTriggerProgrammaticLaunchCompletion(); + + GlmLamportFlags flag(params.bufferFlags, /*numStages=*/1, /*volatileRead=*/false); + auto* stagePtrMcast = static_cast<__nv_bfloat16*>(flag.getCurLamportBuf(params.mcastPtr)); + auto* stagePtrLocal = static_cast<__nv_bfloat16*>(flag.getCurLamportBuf(params.bufferPtrs[params.rank])); + auto& storage = fastEpilogueStorage(); + copyF4(&storage.residual[threadIdx.x], ¶ms.residualIn[threadOffset]); + copyF4(&storage.gamma[threadIdx.x], ¶ms.gamma[packedIdx * kEltsPerThread]); + __pipeline_commit(); + if (isLoader) + { + sPermutedIdx[threadIdx.x] = permutedIdxPart; + sExpertWeightBits[threadIdx.x] = expertWeightBitsPart; + } + __syncthreads(); + + // ==================== Finalize: this thread's elements of the routed sum plus the addend ==================== + int permutedIdx[kMnnvlHopMoeTopK]; +#pragma unroll + for (int k = 0; k < kMnnvlHopMoeTopK; k++) + { + permutedIdx[k] = sPermutedIdx[k]; + } + PackedBf16x8 addend; + addend.packed = *reinterpret_cast(&moe.addend[threadOffset]); + PackedBf16x8 rows[kMnnvlHopMoeTopK]; +#pragma unroll + for (int k = 0; k < kMnnvlHopMoeTopK; k++) + { + int64_t const rowOffset + = static_cast(permutedIdx[k]) * moe.hiddenDimPadded + packedIdx * kEltsPerThread; + rows[k].packed = permutedIdx[k] == -1 ? make_float4(0.F, 0.F, 0.F, 0.F) + : *reinterpret_cast(&moe.fc2Output[rowOffset]); + } + float routed[kEltsPerThread]; +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + routed[i] = 0.F; + } +#pragma unroll + for (int k = 0; k < kMnnvlHopMoeTopK; k++) + { + if (permutedIdx[k] == -1) + { + continue; + } + float const weight = cuda_cast(__ushort_as_bfloat16(sExpertWeightBits[k])); +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + routed[i] = __fmaf_rn(weight, cuda_cast(rows[k].elements[i]), routed[i]); + } + } + PackedBf16x8 val; +#pragma unroll + for (int i = 0; i < kEltsPerThread; i++) + { + val.elements[i] = cuda_cast<__nv_bfloat16, float>(cuda_cast(addend.elements[i]) + + cuda_cast(cuda_cast<__nv_bfloat16, float>(routed[i]))); + } + + // ==================== Broadcast this rank's rows ==================== + sanitizeLamportPayload(val); + reinterpret_cast( + &stagePtrMcast[token * params.tokenDim * WorldSize + params.rank * params.tokenDim])[packedIdx] + = val.packed; + flag.ctaArrive(); + flag.clearDirtyLamportBuf(params.bufferPtrs[params.rank]); + + // ==================== Reduction ==================== + bool rotationPending = threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0; + uint32_t arrivals = 0; + PackedBf16x8 const packedAccum = reduceOneshotPollAll(val, params.rank, stagePtrLocal, token, + params.tokenDim, packedIdx, rotationPending ? flag.arrivalCounter() : nullptr, &arrivals); + uint4 const bytesToClear{ + static_cast(params.numTokens * params.tokenDim * WorldSize * sizeof(__nv_bfloat16)), 0, 0, 0}; + rotateIfAllArrived(flag, rotationPending, arrivals, bytesToClear); + + // ==================== Residual + RMSNorm ==================== + __pipeline_wait_prior(0); + PackedBf16x8 residualIn; + residualIn.packed = storage.residual[threadIdx.x]; + hopEpilogue(params, flag, packedAccum, residualIn, threadOffset, rotationPending, arrivals, bytesToClear); +#endif +} + +// (CTA size, cluster size) of the one-shot grid for rows of tokenDim BF16 elements, eight per thread: the grid of the +// one-shot MNNVL all-reduce (its grid configuration with clusters and one access per scale group). Starts from the +// widest cluster that splits the row evenly, trades cluster width for CTAs of at least 128 threads, narrows the +// clusters while numTokens of them exceed the SMs, and widens them again down to 64-thread CTAs while the SMs allow. +std::pair hopGridConfig(int numTokens, int tokenDim, int smCount) +{ + int const threadsNeeded = (tokenDim + kEltsPerThread - 1) / kEltsPerThread; + int clusterSize = kMaxClusterSize; + while (threadsNeeded % clusterSize != 0 && clusterSize > 1) + { + clusterSize /= 2; + } + int const maxDivisibleClusterSize = clusterSize; + int blockSize = (threadsNeeded + clusterSize - 1) / clusterSize; + while (blockSize < 128 && clusterSize >= 2) + { + blockSize *= 2; + clusterSize /= 2; + } + while (numTokens * clusterSize > smCount && clusterSize > 1 && blockSize <= 512) + { + blockSize *= 2; + clusterSize /= 2; + } + while (clusterSize < maxDivisibleClusterSize) + { + int const candidateClusterSize = clusterSize * 2; + int const candidateBlockSize = (threadsNeeded + candidateClusterSize - 1) / candidateClusterSize; + if (candidateBlockSize < 64 || numTokens * candidateClusterSize > smCount) + { + break; + } + clusterSize = candidateClusterSize; + blockSize = candidateBlockSize; + } + return {blockSize, clusterSize}; +} + +// The launch attributes of both hops: PDL when enabled, the cluster, and the largest shared-memory carveout, so that a +// dependent launched early whose CTAs need most of the shared memory (the fused-A GEMM) can place them on the SMs that +// run this kernel's CTAs instead of starting them only when these exit. +struct HopLaunch +{ + cudaLaunchAttribute attrs[3]; + cudaLaunchConfig_t config; +}; + +void initHopLaunch(HopLaunch& launch, MnnvlHopParams const& params) +{ + TLLM_CHECK_WITH_INFO(common::getSMVersion() >= 90, "[GLM MNNVL hop] requires SM 90 or newer."); + TLLM_CHECK_WITH_INFO(mnnvlHopSupported(params.numTokens, params.tokenDim, params.nRanks), + "[GLM MNNVL hop] unsupported shape: %d tokens, hidden %d, %d ranks.", params.numTokens, params.tokenDim, + params.nRanks); + auto const [blockSize, clusterSize] + = hopGridConfig(params.numTokens, params.tokenDim, common::getMultiProcessorCount()); + launch.attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + launch.attrs[0].val.programmaticStreamSerializationAllowed = common::getEnvEnablePDL() ? 1 : 0; + launch.attrs[1].id = cudaLaunchAttributeClusterDimension; + launch.attrs[1].val.clusterDim.x = 1; + launch.attrs[1].val.clusterDim.y = clusterSize; + launch.attrs[1].val.clusterDim.z = 1; + launch.attrs[2].id = cudaLaunchAttributePreferredSharedMemoryCarveout; + launch.attrs[2].val.sharedMemCarveout = cudaSharedmemCarveoutMaxShared; + launch.config = {}; + launch.config.gridDim = dim3(params.numTokens, clusterSize, 1); + launch.config.blockDim = dim3(blockSize, 1, 1); + launch.config.dynamicSmemBytes = 0; + launch.config.stream = params.stream; + launch.config.attrs = launch.attrs; + launch.config.numAttrs = 3; +} + +HopKernelParams makeHopKernelParams(MnnvlHopParams const& params) +{ + HopKernelParams kernelParams{}; + kernelParams.normOut = static_cast<__nv_bfloat16*>(params.normOut); + kernelParams.residualOut = static_cast<__nv_bfloat16*>(params.residualOut); + kernelParams.shard = static_cast<__nv_bfloat16 const*>(params.input); + kernelParams.residualIn = static_cast<__nv_bfloat16 const*>(params.residualIn); + kernelParams.gamma = static_cast<__nv_bfloat16 const*>(params.gamma); + kernelParams.bufferPtrs = reinterpret_cast<__nv_bfloat16**>(params.bufferPtrsDev); + kernelParams.mcastPtr = static_cast<__nv_bfloat16*>(params.multicastPtr); + kernelParams.bufferFlags = params.bufferFlags; + kernelParams.numTokens = params.numTokens; + kernelParams.tokenDim = params.tokenDim; + kernelParams.rank = params.rank; + kernelParams.epsilon = params.epsilon; + return kernelParams; +} + +template +void launchHop(cudaLaunchConfig_t const& config, HopKernelParams const& kernelParams, bool shardPublished) +{ + if (shardPublished) + { + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, &mnnvlHopKernel, kernelParams)); + } + else + { + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, &mnnvlHopKernel, kernelParams)); + } +} + +template +void launchMoeFinalizeHop( + cudaLaunchConfig_t const& config, HopKernelParams const& kernelParams, MoeFinalizeKernelParams const& moeParams) +{ + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, &mnnvlMoeFinalizeHopKernel, kernelParams, moeParams)); +} + +} // namespace + +int64_t mnnvlHopBufferBytes(int numTokens, int tokenDim, int nRanks) +{ + return static_cast(numTokens) * tokenDim * nRanks * static_cast(sizeof(__nv_bfloat16)); +} + +bool mnnvlHopSupported(int numTokens, int tokenDim, int nRanks) +{ + if ((nRanks != 2 && nRanks != 4 && nRanks != 8 && nRanks != 16) || numTokens < 1 || tokenDim <= 0 + || tokenDim % kEltsPerThread != 0) + { + return false; + } + int64_t const bytes = mnnvlHopBufferBytes(numTokens, tokenDim, nRanks); + if (bytes > kOneShotBytes && !(nRanks == 16 && bytes <= kOneShotBytes16Ranks)) + { + return false; + } + auto const [blockSize, clusterSize] = hopGridConfig(numTokens, tokenDim, common::getMultiProcessorCount()); + return blockSize % kWarpSize == 0 && blockSize <= kMaxThreads && clusterSize > 1 && clusterSize <= kMaxClusterSize + && blockSize * clusterSize * kEltsPerThread == tokenDim; +} + +void mnnvlHop(MnnvlHopParams const& params) +{ + HopLaunch launch; + initHopLaunch(launch, params); + HopKernelParams const kernelParams = makeHopKernelParams(params); + switch (params.nRanks) + { + case 2: launchHop<2>(launch.config, kernelParams, params.shardPublished); break; + case 4: launchHop<4>(launch.config, kernelParams, params.shardPublished); break; + case 8: launchHop<8>(launch.config, kernelParams, params.shardPublished); break; + case 16: launchHop<16>(launch.config, kernelParams, params.shardPublished); break; + default: TLLM_CHECK_WITH_INFO(false, "[GLM MNNVL hop] unsupported world size %d.", params.nRanks); + } +} + +void mnnvlMoeFinalizeHop(MnnvlHopParams const& params, MnnvlMoeFinalizeParams const& moe) +{ + HopLaunch launch; + initHopLaunch(launch, params); + HopKernelParams const kernelParams = makeHopKernelParams(params); + MoeFinalizeKernelParams moeParams{}; + moeParams.fc2Output = static_cast<__nv_bfloat16 const*>(moe.fc2Output); + moeParams.expertWeights = static_cast<__nv_bfloat16 const*>(moe.expertWeights); + moeParams.expandedIdxToPermutedIdx = moe.expandedIdxToPermutedIdx; + moeParams.addend = static_cast<__nv_bfloat16 const*>(moe.addend); + moeParams.hiddenDimPadded = moe.hiddenDimPadded; + switch (params.nRanks) + { + case 2: launchMoeFinalizeHop<2>(launch.config, kernelParams, moeParams); break; + case 4: launchMoeFinalizeHop<4>(launch.config, kernelParams, moeParams); break; + case 8: launchMoeFinalizeHop<8>(launch.config, kernelParams, moeParams); break; + case 16: launchMoeFinalizeHop<16>(launch.config, kernelParams, moeParams); break; + default: TLLM_CHECK_WITH_INFO(false, "[GLM MNNVL MoE-finalize hop] unsupported world size %d.", params.nRanks); + } +} + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.h new file mode 100644 index 000000000000..aaa7d02be3af --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.h @@ -0,0 +1,99 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +//! Routed experts per token of the MoE-finalize hop. +constexpr int kMnnvlHopMoeTopK = 8; + +//! \brief One BF16 one-shot MNNVL hop: the all-reduce of [numTokens, tokenDim] rows over nRanks ranks, plus the +//! residual and RMSNorm, on an MNNVL all-reduce workspace (three Lamport buffers and their flags, -0.0 as the empty +//! value; the workspace of the one-shot MNNVL all-reduce). +struct MnnvlHopParams +{ + int nRanks; //!< Ranks of the group: 2, 4, 8 or 16 + int rank; //!< This rank + int numTokens; //!< Rows + int tokenDim; //!< Elements per row + void** bufferPtrsDev; //!< Unicast addresses of every rank's workspace (a device array) + void* multicastPtr; //!< Multicast address of the workspace + uint32_t* bufferFlags; //!< Lamport flags of the workspace + void const* input; //!< This rank's partial rows [numTokens, tokenDim]; not read when shardPublished + void const* residualIn; //!< [numTokens, tokenDim] + void const* gamma; //!< RMSNorm weight [tokenDim] + float epsilon; //!< RMSNorm epsilon + void* residualOut; //!< Rows summed over the ranks plus residualIn [numTokens, tokenDim] + void* normOut; //!< RMSNorm of residualOut times gamma [numTokens, tokenDim] + bool shardPublished; //!< The producer of input already stored this rank's rows into the workspace + cudaStream_t stream; +}; + +//! \brief The routed experts' rows of the MoE-finalize hop: the do_finalize=false outputs of the MoE and the +//! shared-expert output added to them. +struct MnnvlMoeFinalizeParams +{ + void const* fc2Output; //!< FC2 rows [*, hiddenDimPadded], BF16 + int hiddenDimPadded; //!< Row pitch of fc2Output in elements + void const* expertWeights; //!< Routing weights [numTokens, kMnnvlHopMoeTopK], BF16 + int32_t const* expandedIdxToPermutedIdx; //!< FC2 row of each (token, k) [numTokens, kMnnvlHopMoeTopK]; -1 skips + void const* addend; //!< Added to the finalized rows [numTokens, tokenDim], BF16 +}; + +//! Whether the hops serve a BF16 [numTokens, tokenDim] all-reduce over nRanks ranks: 2, 4, 8 or 16 ranks, one-shot +//! (up to 1 MiB of Lamport data per rank, and up to 1.5 MiB on 16 ranks), and a grid of one row per cluster of 2 to 8 +//! CTAs of at most 128 threads (whole warps), eight elements per thread, that covers the row exactly. +bool mnnvlHopSupported(int numTokens, int tokenDim, int nRanks); + +//! Bytes of one Lamport buffer that a hop of numTokens x tokenDim BF16 elements over nRanks ranks uses. +int64_t mnnvlHopBufferBytes(int numTokens, int tokenDim, int nRanks); + +//! \brief The one-shot all-reduce with the residual add and RMSNorm (FastEpilogue, FastProtocol, one poll loop). +//! +//! Every rank stores its rows (-0.0 written as +0.0) through the multicast address at [token][rank][column] of the +//! current Lamport buffer, empties the dirty buffer, and polls all slots in one loop. The sum is in rank order in FP32, +//! from zero per chunk of 8 ranks, the chunks added in order, and rounded to BF16; the residual is added in BF16; the +//! sum of squares of the BF16 values goes through a transaction barrier of the cluster. The PDL dependents launch +//! right after the grid dependency wait, so they must read the outputs only after their own wait. +//! +//! With shardPublished the producer of input (the o_proj GEMM's Lamport publish) has stored this rank's rows into +//! every rank's current buffer: the kernel polls every slot, its own included, does not wait for its primary before +//! the exchange and reads the flags and the residual with volatile loads; it counts its CTA arrival after its poll +//! and waits for its primary at its end. The producer must trigger its dependents only after its own grid dependency +//! wait. +void mnnvlHop(MnnvlHopParams const& params); + +//! \brief The MoE finalize fused into the hop: each thread finalizes its eight elements of every routed row (the +//! token's FC2 rows weighted in k order with an FMA, -1 slots skipped, rounded to BF16, the addend added in FP32 and +//! the sum rounded once more) and publishes them as its rows; the rest is mnnvlHop without a published shard. +//! params.input and params.shardPublished are not read. Eight threads per CTA read the token's FC2 row indices and +//! routing weights before the grid dependency wait: the kernel that wrote them must have completed when this grid +//! launches (the MoE's FC2 triggers its dependents only after its own wait). +void mnnvlMoeFinalizeHop(MnnvlHopParams const& params, MnnvlMoeFinalizeParams const& moe); + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.cu new file mode 100644 index 000000000000..1b5a83166840 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.cu @@ -0,0 +1,244 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +namespace +{ + +// Convert 8 bfloat16 values from a uint4 to float array - optimized conversion +template +__device__ __forceinline__ void bf16_uint4_to_float8(uint4 const& vec, float* dst) +{ + __nv_bfloat16* bf16_ptr = reinterpret_cast<__nv_bfloat16*>(const_cast(&vec)); + +#pragma unroll + for (int i = 0; i < VPT; i++) + { + dst[i] = __bfloat162float(bf16_ptr[i]); + } +} + +// mat_a: the activations [kNumTokens, kHiddenDim]; mat_b: the weight [kNumExperts, kHiddenDim]; out: the logits +// [kNumTokens, kNumExperts]. Before the grid-dependency wait each CTA prefetches its weight row into L2 (no kernel of +// the step writes the weight), so the weight loads after the wait hit L2; the dependent grid may launch right after the +// wait instead of after the logits are written, and must wait on its own grid dependency before reading them. +template +__global__ __launch_bounds__(128, 1) void router_gemm_kernel(float* out, T const* mat_a, T const* mat_b) +{ + // Each block handles one expert column + int const n_idx = blockIdx.x; + int const tid = threadIdx.x; + constexpr int kWarpSize = 32; + constexpr int kNumWarps = kBlockSize / kWarpSize; + // Constants for this kernel + constexpr int k_elems_per_k_iteration = VPT * kBlockSize; + constexpr int k_iterations = kHiddenDim / k_elems_per_k_iteration; // Total K iterations + + // Initialize accumulators for all M rows + float acc[kNumTokens] = {}; + + // Shared memory for warp-level reduction + __shared__ float sm_reduction[kNumTokens][kNumWarps]; // kNumWarps + + // The weight row of expert n_idx (a column of B in column-major order) + T const* b_col = mat_b + n_idx * kHiddenDim; + + // Pre-compute k_base values for each iteration to help compiler optimize + int k_bases[k_iterations]; +#pragma unroll + for (int ki = 0; ki < k_iterations; ki++) + { + k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT; + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + { + // The primary grid does not write the weights, so this CTA's column can stream into L2 while the primary + // runs; the weight loads below then hit L2. What is loaded and the accumulation order are unchanged. + constexpr int kLineBytes = 128; + constexpr int kColumnBytes = kHiddenDim * static_cast(sizeof(T)); + static_assert(kColumnBytes % kLineBytes == 0, "the weight column must span whole 128-byte lines"); + char const* b_col_bytes = reinterpret_cast(b_col); + for (int offset = tid * kLineBytes; offset < kColumnBytes; offset += kBlockSize * kLineBytes) + { + asm volatile("prefetch.global.L2 [%0];" ::"l"(b_col_bytes + offset) : "memory"); + } + } + cudaGridDependencySynchronize(); + // Every dependent waits on its own grid dependency before it reads the logits, so it may launch now and overlap + // its prologue with this GEMV. + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + // Process the GEMM in chunks + for (int ki = 0; ki < k_iterations; ki++) + { + int const k_base = k_bases[ki]; + + // Load B matrix values using vector load (8 bf16 values) + uint4 b_vec = *reinterpret_cast(b_col + k_base); + + // Convert B values to float + float b_float[VPT]; + bf16_uint4_to_float8(b_vec, b_float); + +// Process each token +#pragma unroll + for (int m_idx = 0; m_idx < kNumTokens; m_idx++) + { + // Load both rows of A matrix using vector loads + uint4 a_vec = *reinterpret_cast(mat_a + (m_idx * kHiddenDim) + k_base); + + // Convert A values to float + float a_float[VPT]; + bf16_uint4_to_float8(a_vec, a_float); + +// Process elements in this chunk +#pragma unroll + for (int k = 0; k < VPT; k++) + { + float a = a_float[k]; + float b = b_float[k]; + acc[m_idx] += a * b; + } + } + } + + // Perform warp-level reduction + int const warpSize = 32; + int const warpId = tid / warpSize; + int const laneId = tid % warpSize; + + // Register for warp-level reduction results + float warp_result[kNumTokens]; + +#pragma unroll + for (int m_idx = 0; m_idx < kNumTokens; m_idx++) + { + warp_result[m_idx] = acc[m_idx]; + } + +// Perform warp-level reduction using optimized butterfly pattern +#pragma unroll + for (int m = 0; m < kNumTokens; m++) + { + float sum = warp_result[m]; + + // Butterfly reduction pattern + sum += __shfl_xor_sync(0xffffffff, sum, 16); + sum += __shfl_xor_sync(0xffffffff, sum, 8); + sum += __shfl_xor_sync(0xffffffff, sum, 4); + sum += __shfl_xor_sync(0xffffffff, sum, 2); + sum += __shfl_xor_sync(0xffffffff, sum, 1); + + // Only the first thread in each warp stores to shared memory + if (laneId == 0) + { + sm_reduction[m][warpId] = sum; + } + } + + __syncthreads(); + + // Final reduction across warps (only first thread) + if (tid == 0) + { +#pragma unroll + for (int m = 0; m < kNumTokens; m++) + { + float final_sum = 0.0f; + +// Sum across the kNumWarps +#pragma unroll + for (int w = 0; w < kNumWarps; w++) + { + final_sum += sm_reduction[m][w]; + } + + // Write final result + out[m * kNumExperts + n_idx] = final_sum; + } + } +} + +template +void launchRouterGemm( + float* logits, __nv_bfloat16 const* hiddenStates, __nv_bfloat16 const* weight, cudaStream_t stream) +{ + using T = __nv_bfloat16; + constexpr int VPT = 16 / sizeof(T); + constexpr int kBlockSize = 128; + cudaLaunchConfig_t config; + config.gridDim = kRouterNumExperts; + config.blockDim = kBlockSize; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = attrs; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, + router_gemm_kernel, logits, hiddenStates, + weight)); +} + +// Launches the instance of numTokens, one of kNumTokens..kRouterMaxTokens. +template +void dispatchRouterGemm( + float* logits, __nv_bfloat16 const* hiddenStates, __nv_bfloat16 const* weight, int numTokens, cudaStream_t stream) +{ + if constexpr (kNumTokens <= kRouterMaxTokens) + { + if (numTokens == kNumTokens) + { + launchRouterGemm(logits, hiddenStates, weight, stream); + } + else + { + dispatchRouterGemm(logits, hiddenStates, weight, numTokens, stream); + } + } + else + { + TLLM_THROW("The GLM router GEMV takes 1 to %d tokens, got %d", kRouterMaxTokens, numTokens); + } +} + +} // namespace + +void invokeRouterGemm( + float* logits, __nv_bfloat16 const* hiddenStates, __nv_bfloat16 const* weight, int numTokens, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO( + numTokens >= 1, "The GLM router GEMV takes 1 to %d tokens, got %d", kRouterMaxTokens, numTokens); + dispatchRouterGemm<1>(logits, hiddenStates, weight, numTokens, stream); +} + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.h new file mode 100644 index 000000000000..c103c7eeb7e7 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.h @@ -0,0 +1,50 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +//! The router GEMV's K: the GLM-5 hidden size. +constexpr int kRouterHiddenSize = 6144; +//! Routed experts: one CTA computes one expert's logits. +constexpr int kRouterNumExperts = 256; +//! Tokens per call at most: one kernel instance per token count. +constexpr int kRouterMaxTokens = 16; + +//! logits[t, e] = sum_k hiddenStates[t, k] * weight[e, k]: BF16 operands, FP32 accumulation and FP32 logits, all +//! dense row-major. hiddenStates is [numTokens, 6144] with numTokens in [1, 16], weight [256, 6144] and logits +//! [numTokens, 256]. +//! +//! Launched with the programmatic-dependent-launch attribute when TRTLLM_ENABLE_PDL is set. Before its grid-dependency +//! wait each CTA prefetches its weight row into L2; right after the wait, before any logit is written, it triggers its +//! dependents. Every kernel that follows on the stream must therefore wait on its own grid dependency before it reads +//! the logits or writes hiddenStates. +void invokeRouterGemm( + float* logits, __nv_bfloat16 const* hiddenStates, __nv_bfloat16 const* weight, int numTokens, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.cu new file mode 100644 index 000000000000..af9aade1887d --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.cu @@ -0,0 +1,817 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ +namespace +{ + +// The MMAs run 8 head columns: the kHeads query heads in columns 0 .. kHeads - 1, zero Q rows in the others. +constexpr int kMmaHeads = 8; +constexpr int kHeads = kGlmSparseMlaNumHeads; +constexpr int kHeadDimQk = kGlmSparseMlaHeadDimQk; +constexpr int kHeadDimV = kGlmSparseMlaHeadDimV; +// One cluster of kSplits CTAs per query token; CTA s owns indices [s * kTokensPerCta, (s + 1) * kTokensPerCta). +constexpr int kSplits = 16; +constexpr int kTokensPerCta = kGlmSparseMlaMaxTopK / kSplits; +constexpr int kThreads = 256; +constexpr int kWarps = kThreads / 32; +constexpr int kTokensPerWarp = kTokensPerCta / kWarps; +static_assert(kTokensPerWarp == 16, "a warp gathers one 16-token MMA tile, one token row per lane"); +static_assert(kHeads < kMmaHeads, "the query heads fill part of the MMA head columns"); +// A token row: 576 E4M3 values. In shared memory the rows are 592 bytes apart, so the eight 16-byte rows of every +// ldmatrix phase fall into distinct banks. +constexpr int kRowBytes = kHeadDimQk; +constexpr int kSmemRowBytes = kRowBytes + 16; +// P (E4M3) per MMA head column, padded the same way. +constexpr int kPRowBytes = kTokensPerCta + 16; +// Output columns merged by each CTA of the cluster. +constexpr int kDvPerCta = kHeadDimV / kSplits; +static_assert(kDvPerCta % 16 == 0, "each output slice holds whole 16-column MMA blocks"); +// P is scaled to the E4M3 range before the second MMA; the normalization divides it out. +constexpr float kPScale = 448.f; + +// v_b BMM: CTA r of the cluster owns v in [r * kVPerCta, (r + 1) * kVPerCta) of every head. +constexpr int kVHeadDim = kGlmSparseMlaVHeadDim; +constexpr int kVPerCta = kVHeadDim / kSplits; +static_assert(kVPerCta == 16, "a warp computes one head's 16 rows of the v_b slice as one m16n8k16 row tile"); +// The v_b_proj rows [h][r * 16, r * 16 + 16)[0, 512) of one head: 16 KiB, contiguous in global memory. +constexpr int kVbHeadBytes = kVPerCta * kHeadDimV * 2; + +// Shared memory: token rows, Q rows, P, the per-warp max / sum, the receive buffers of the cluster merge, the v_b +// slice, the CTA's merged output slice, the row's gathered BF16 output and the mbarriers. +constexpr int kOffQ = kTokensPerCta * kSmemRowBytes; +constexpr int kOffP = kOffQ + kMmaHeads * kSmemRowBytes; +constexpr int kOffRed = kOffP + kMmaHeads * kPRowBytes; +constexpr int kRedFloats = 2 * kWarps * kMmaHeads; +constexpr int kOffRecvO = kOffRed + kRedFloats * 4; +constexpr int kRecvOFloats = kSplits * kMmaHeads * kDvPerCta; +constexpr int kOffRecvStats = kOffRecvO + kRecvOFloats * 4; +constexpr int kRecvStatsFloats = kSplits * kMmaHeads * 2; +constexpr int kRecvBytes = (kRecvOFloats + kRecvStatsFloats) * 4; +constexpr int kOffVb = kOffRecvStats + kRecvStatsFloats * 4; +constexpr int kVbBytes = kHeads * kVbHeadBytes; +// The CTA's merged output slice in BF16, [8 head columns][32 columns]. +constexpr int kOffSlice = kOffVb + kVbBytes; +constexpr int kSliceBytes = kMmaHeads * kDvPerCta * 2; +// The row's merged output in BF16, [8 head columns][512]. +constexpr int kOffGatheredO = kOffSlice + kSliceBytes; +constexpr int kGatheredOBytes = kMmaHeads * kHeadDimV * 2; +constexpr int kOffBars = kOffGatheredO + kGatheredOBytes; +// One gather barrier per warp, the Q barrier, the receive barrier, the v_b slice barrier and the gathered-output +// barrier. +constexpr int kNumBars = kWarps + 4; +constexpr int kSmemBytes = kOffBars + kNumBars * 8; + +static_assert(kOffQ % 16 == 0 && kOffP % 16 == 0 && kOffRecvO % 16 == 0 && kOffVb % 16 == 0 && kOffSlice % 16 == 0 + && kOffGatheredO % 16 == 0 && kOffBars % 8 == 0, + "alignment"); +static_assert(kSplits * kMmaHeads <= kThreads, "one max / sum transfer per thread"); +static_assert(kMmaHeads * kDvPerCta == kThreads, "the merge maps one output element to each thread"); +static_assert(kDvPerCta == 32 && kSplits <= 32, "the merge maps warp h to head h of the slice, lane r to partial r"); +static_assert(kSmemBytes <= 227 * 1024, "shared memory per CTA"); + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + +__device__ __forceinline__ uint32_t smemAddress(void const* ptr) +{ + return static_cast(__cvta_generic_to_shared(ptr)); +} + +__device__ __forceinline__ void mbarrierInit(uint32_t bar, uint32_t count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(bar), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrierArriveExpectTx(uint32_t bar, uint32_t bytes) +{ + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(bar), "r"(bytes) : "memory"); +} + +__device__ __forceinline__ void mbarrierWait(uint32_t bar, uint32_t parity) +{ + asm volatile( + "{\n" + ".reg .pred complete;\n" + "GLM_SPARSE_MLA_FD_WAIT_%=:\n" + "mbarrier.try_wait.parity.shared::cta.b64 complete, [%0], %1;\n" + "@!complete bra GLM_SPARSE_MLA_FD_WAIT_%=;\n" + "}\n" ::"r"(bar), + "r"(parity) + : "memory"); +} + +// Wait with cluster-scope acquire: the receive barriers complete on bytes that other CTAs of the cluster wrote. +__device__ __forceinline__ void mbarrierWaitCluster(uint32_t bar, uint32_t parity) +{ + asm volatile( + "{\n" + ".reg .pred complete;\n" + "GLM_SPARSE_MLA_FD_WAIT_CLUSTER_%=:\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 complete, [%0], %1;\n" + "@!complete bra GLM_SPARSE_MLA_FD_WAIT_CLUSTER_%=;\n" + "}\n" ::"r"(bar), + "r"(parity) + : "memory"); +} + +__device__ __forceinline__ void bulkCopyGlobalToShared(uint32_t dst, void const* src, uint32_t bytes, uint32_t bar) +{ + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n" ::"r"(dst), + "l"(src), "r"(bytes), "r"(bar) + : "memory"); +} + +// Two floats into the shared memory of a CTA of the cluster; the bytes complete on that CTA's barrier. +__device__ __forceinline__ void storeRemote(uint32_t remoteAddr, float a, float b, uint32_t remoteBar) +{ + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.f32 [%0], {%1, %2}, [%3];\n" ::"r"(remoteAddr), + "f"(a), "f"(b), "r"(remoteBar) + : "memory"); +} + +// Sixteen bytes into the shared memory of a CTA of the cluster; the bytes complete on that CTA's barrier. +__device__ __forceinline__ void storeRemote16(uint32_t remoteAddr, uint4 value, uint32_t remoteBar) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];\n" ::"r"( + remoteAddr), + "r"(value.x), "r"(value.y), "r"(value.z), "r"(value.w), "r"(remoteBar) + : "memory"); +} + +__device__ __forceinline__ uint32_t mapToRank(uint32_t addr, uint32_t rank) +{ + uint32_t mapped; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;\n" : "=r"(mapped) : "r"(addr), "r"(rank)); + return mapped; +} + +__device__ __forceinline__ void clusterArriveRelaxed() +{ + asm volatile("barrier.cluster.arrive.relaxed.aligned;\n" ::: "memory"); +} + +__device__ __forceinline__ void clusterWait() +{ + asm volatile("barrier.cluster.wait.acquire.aligned;\n" ::: "memory"); +} + +__device__ __forceinline__ void ldmatrixX4(uint32_t addr, uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) + : "r"(addr)); +} + +__device__ __forceinline__ void ldmatrixX4Trans(uint32_t addr, uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) + : "r"(addr)); +} + +__device__ __forceinline__ void ldmatrixX2(uint32_t addr, uint32_t& r0, uint32_t& r1) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" : "=r"(r0), "=r"(r1) : "r"(addr)); +} + +// D += A B with A 16 x 16 (row major) and B 16 x 8 (column major), BF16 inputs and fp32 accumulation. +__device__ __forceinline__ void mmaBf16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +// D += A B with A 16 x 16 (row major) and B 16 x 8 (column major), FP16 inputs and fp32 accumulation. +__device__ __forceinline__ void mmaF16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +// The two E4M3 values in the low 16 bits of `pair` as two FP16 values (exact). +__device__ __forceinline__ uint32_t e4m3x2ToHalf2(uint32_t pair) +{ + __half2_raw const h = __nv_cvt_fp8x2_to_halfraw2(static_cast<__nv_fp8x2_storage_t>(pair & 0xffffU), __NV_E4M3); + return static_cast(h.x) | (static_cast(h.y) << 16); +} + +// The B fragment of an E4M3 m16n8k32 MMA ({b0, b1}) unpacked for the two FP16 MMAs of mmaE4m3AsF16: low halves, then +// high halves. +__device__ __forceinline__ void unpackE4m3B(uint32_t (&bh)[4], uint32_t b0, uint32_t b1) +{ + bh[0] = e4m3x2ToHalf2(b0); + bh[1] = e4m3x2ToHalf2(b1); + bh[2] = e4m3x2ToHalf2(b0 >> 16); + bh[3] = e4m3x2ToHalf2(b1 >> 16); +} + +// D += A B for an E4M3 m16n8k32 MMA with A in its E4M3 fragment registers and B unpacked beforehand (unpackE4m3B): +// an FP16 m16n8k16 MMA on the low 16-bit halves of the A registers, then one on the high halves, chained on D (the +// SM100 family has no E4M3 mma.sync; this is the sequence the compiler lowers one to). +__device__ __forceinline__ void mmaE4m3AsF16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t const (&bh)[4]) +{ + mmaF16(d, e4m3x2ToHalf2(a0), e4m3x2ToHalf2(a1), e4m3x2ToHalf2(a2), e4m3x2ToHalf2(a3), bh[0], bh[1]); + mmaF16(d, e4m3x2ToHalf2(a0 >> 16), e4m3x2ToHalf2(a1 >> 16), e4m3x2ToHalf2(a2 >> 16), e4m3x2ToHalf2(a3 >> 16), bh[2], + bh[3]); +} + +__device__ __forceinline__ uint8_t toE4m3(float x) +{ + return static_cast(__nv_cvt_float_to_fp8(x, __NV_SATFINITE, __NV_E4M3)); +} + +// Position of the CTA's token t in the k order of the P.V MMA: within every 16 tokens, k slots 4c .. 4c + 3 hold the +// tokens {2c, 2c + 1, 2c + 8, 2c + 9}, the order in which ldmatrix.trans of the token-major V rows plus two byte +// permutes deliver a V^T fragment. +__device__ __forceinline__ int permutedSlot(int t) +{ + int const u = t & 15; + int const slot16 = (u < 8) ? ((u >> 1) * 4 + (u & 1)) : (((u - 8) >> 1) * 4 + 2 + (u & 1)); + return (t & ~15) + slot16; +} + +// The row of one lane of a warp's gather, found before the grid-dependency wait (kGatherBeforeWait). +struct EarlyGatherRow +{ + int64_t row; + bool valid; + // A position of this step's tokens: its row is written by the KV append this grid depends on. + bool fresh; +}; + +// Lane l < 16 of the warp maps index slot `slot` of query token `tokenQ` (a position) to its pool row with the remap +// arithmetic, and bulk-copies the row unless it holds one of this step's tokens. Masked lanes copy row 0, as after the +// wait. Every read here (indices, block table, KV length, rows of earlier steps) precedes this grid's launch. +__device__ __forceinline__ EarlyGatherRow gatherRowBeforeWait( + GlmSparseMlaFlashDecodeParams const& params, int tokenQ, int slot, int lane, uint32_t dst, uint32_t barGather) +{ + int position = -1; + if (lane < kTokensPerWarp && slot < params.topK) + { + position = __ldg(params.indices + static_cast(tokenQ) * params.indicesStride + slot); + } + int const seqIdx = tokenQ / params.seqLenQ; + int const tokenInSeq = tokenQ - seqIdx * params.seqLenQ; + int const block = position / params.remapTokensPerBlock; + int page = -1; + if (position >= 0 && block < params.remapMaxBlocks) + { + page = __ldg(params.remapBlockTable + seqIdx * params.remapBlockTableStride + block); + } + int const inBlock = position - block * params.remapTokensPerBlock; + int const rowRaw = page >= 0 ? page * params.remapStrideFactor + inBlock + params.remapLayerOffset : -1; + int const seqLenKv = __ldg(params.seqLensKv + seqIdx); + int const numValid = min(seqLenKv - (params.seqLenQ - 1 - tokenInSeq), params.topK); + EarlyGatherRow out; + out.valid = lane < kTokensPerWarp && slot < numValid && rowRaw >= 0; + out.row = out.valid ? rowRaw : 0; + out.fresh = out.valid && position >= seqLenKv - params.seqLenQ; + if (lane == 0) + { + mbarrierArriveExpectTx(barGather, kTokensPerWarp * kRowBytes); + } + __syncwarp(); + if (lane < kTokensPerWarp && !out.fresh) + { + bulkCopyGlobalToShared( + dst, static_cast(params.kvPool) + out.row * kRowBytes, kRowBytes, barGather); + } + return out; +} + +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + +// kRemap: the indices are positions, mapped to pool rows through the block table. kGatherBeforeWait (with kRemap): the +// indices, the block table and the KV length are read, and the rows of the positions before this step's tokens +// gathered, before the grid-dependency wait; the rows of this step's tokens after it. +template +__global__ void __launch_bounds__(kThreads, 1) glmSparseMlaFlashDecodeKernel(GlmSparseMlaFlashDecodeParams const params) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + static_assert(kRemap || !kGatherBeforeWait, "the early gather tells this step's tokens by their positions"); + extern __shared__ __align__(1024) uint8_t smem[]; + uint8_t* sKv = smem; + uint8_t* sQ = smem + kOffQ; + uint8_t* sP = smem + kOffP; + float* sRedMax = reinterpret_cast(smem + kOffRed); + float* sRedSum = sRedMax + kWarps * kMmaHeads; + float* recvO = reinterpret_cast(smem + kOffRecvO); + float* recvStats = reinterpret_cast(smem + kOffRecvStats); + uint8_t* sSlice = smem + kOffSlice; + uint8_t* sGatheredO = smem + kOffGatheredO; + uint64_t* bars = reinterpret_cast(smem + kOffBars); + + int const tid = threadIdx.x; + int const warp = tid / 32; + int const lane = tid % 32; + // The cluster spans gridDim.x, so the split index is the CTA's rank in the cluster. + int const split = blockIdx.x; + int const tokenQ = blockIdx.y; + + uint32_t const barGather = smemAddress(bars + warp); + uint32_t const barQ = smemAddress(bars + kWarps); + uint32_t const barRecv = smemAddress(bars + kWarps + 1); + uint32_t const barVb = smemAddress(bars + kWarps + 2); + uint32_t const barGatheredO = smemAddress(bars + kWarps + 3); + { + // Q rows kHeads .. 7 feed the MMAs' unused head columns: zero, before the barrier below orders them for the QK. + constexpr int kZeroChunks = (kMmaHeads - kHeads) * kRowBytes / 16; + static_assert(kZeroChunks <= kThreads && kRowBytes % 16 == 0, "one 16-byte zero store per thread"); + if (tid < kZeroChunks) + { + int const r = kHeads + tid / (kRowBytes / 16); + int const chunk = tid % (kRowBytes / 16); + *reinterpret_cast(sQ + r * kSmemRowBytes + 16 * chunk) = make_uint4(0, 0, 0, 0); + } + } + if (tid == 0) + { + for (int i = 0; i < kNumBars; ++i) + { + mbarrierInit(smemAddress(bars + i), 1); + } + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + mbarrierArriveExpectTx(barRecv, kRecvBytes); + mbarrierArriveExpectTx(barGatheredO, kGatheredOBytes); + } + __syncthreads(); + // Paired with the wait before the first remote store: no CTA writes into a receive buffer whose barrier other + // CTAs of the cluster have not initialized yet. + clusterArriveRelaxed(); + + [[maybe_unused]] EarlyGatherRow early{}; + if constexpr (kGatherBeforeWait) + { + if (warp == 0) + { + // The block-table row of the token's sequence into L2 while the indices load: the remap below reads it + // next. The CTAs of a cluster prefetch interleaved 128-byte lines of the row. + char const* tableRow = reinterpret_cast( + params.remapBlockTable + (tokenQ / params.seqLenQ) * params.remapBlockTableStride); + int64_t const rowBytes = static_cast(params.remapMaxBlocks) * sizeof(int32_t); + for (int64_t off = (static_cast(split) * 32 + lane) * 128; off < rowBytes; + off += kSplits * 32 * 128) + { + asm volatile("prefetch.global.L2 [%0];\n" ::"l"(tableRow + off)); + } + } + int const slotEarly = split * kTokensPerCta + warp * kTokensPerWarp + lane; + early = gatherRowBeforeWait(params, tokenQ, slotEarly, lane, + smemAddress(sKv + (warp * kTokensPerWarp + lane) * kSmemRowBytes), barGather); + } + + // Q, the indices, the KV pool rows and the scales come from earlier kernels. + cudaGridDependencySynchronize(); + + if (warp == 0) + { + if (lane == 0) + { + mbarrierArriveExpectTx(barQ, kHeads * kRowBytes); + } + __syncwarp(); + if (lane < kHeads) + { + bulkCopyGlobalToShared(smemAddress(sQ + lane * kSmemRowBytes), + static_cast(params.q) + (static_cast(tokenQ) * kHeads + lane) * kRowBytes, + kRowBytes, barQ); + } + } + + // Lane l of warp w gathers token tokBase + l of the CTA. The index is read before the KV length is known, so the + // two loads share one round trip. + int const tokBase = warp * kTokensPerWarp; + uint32_t validMask; + if constexpr (kGatherBeforeWait) + { + // The rows of this step's tokens, which the KV append wrote. + if (lane < kTokensPerWarp && early.fresh) + { + bulkCopyGlobalToShared(smemAddress(sKv + (tokBase + lane) * kSmemRowBytes), + static_cast(params.kvPool) + early.row * kRowBytes, kRowBytes, barGather); + } + validMask = __ballot_sync(0xffffffffU, early.valid); + } + else + { + int const slot = split * kTokensPerCta + tokBase + lane; + int indexRaw = -1; + if (lane < kTokensPerWarp && slot < params.topK) + { + indexRaw = __ldg(params.indices + static_cast(tokenQ) * params.indicesStride + slot); + } + int const seqIdx = tokenQ / params.seqLenQ; + int const tokenInSeq = tokenQ - seqIdx * params.seqLenQ; + if constexpr (kRemap) + { + // The index is a position in the sequence: map it to its pool row through the block table. + int const block = indexRaw / params.remapTokensPerBlock; + int page = -1; + if (indexRaw >= 0 && block < params.remapMaxBlocks) + { + page = __ldg(params.remapBlockTable + seqIdx * params.remapBlockTableStride + block); + } + int const inBlock = indexRaw - block * params.remapTokensPerBlock; + indexRaw = page >= 0 ? page * params.remapStrideFactor + inBlock + params.remapLayerOffset : -1; + } + int const numValid = min(__ldg(params.seqLensKv + seqIdx) - (params.seqLenQ - 1 - tokenInSeq), params.topK); + bool const valid = lane < kTokensPerWarp && slot < numValid && indexRaw >= 0; + // Masked lanes still load a row (row 0 of the pool) so that every gather barrier expects the same bytes. + int64_t const row = valid ? indexRaw : 0; + if (lane == 0) + { + mbarrierArriveExpectTx(barGather, kTokensPerWarp * kRowBytes); + } + __syncwarp(); + if (lane < kTokensPerWarp) + { + bulkCopyGlobalToShared(smemAddress(sKv + (tokBase + lane) * kSmemRowBytes), + static_cast(params.kvPool) + row * kRowBytes, kRowBytes, barGather); + } + validMask = __ballot_sync(0xffffffffU, valid); + } + float const scaleSoftmaxLog2 = __ldg(params.scaleSoftmaxLog2); + float const outputScale = __ldg(params.outputScale); + + // Fragment coordinates of mma.sync: g = row group, c = column pair. + int const g = lane / 4; + int const c = lane % 4; + + // S^T = K Q^T for the warp's 16 tokens and the 8 head columns; two accumulators over alternating k steps. The + // query fragments of every k step are unpacked once, while the warp's token rows are in flight. + mbarrierWait(barQ, 0); + uint32_t qh[kHeadDimQk / 32][4]; + { + uint32_t const qAddr = smemAddress(sQ + (lane % 8) * kSmemRowBytes + 16 * ((lane / 8) % 2)); +#pragma unroll + for (int kd = 0; kd < kHeadDimQk / 32; ++kd) + { + uint32_t b0, b1; + ldmatrixX2(qAddr + 32 * kd, b0, b1); + unpackE4m3B(qh[kd], b0, b1); + } + } + mbarrierWait(barGather, 0); + float s[4]; + { + float acc[2][4] = {{0.f, 0.f, 0.f, 0.f}, {0.f, 0.f, 0.f, 0.f}}; + uint32_t const kAddr + = smemAddress(sKv + (tokBase + lane % 8 + 8 * ((lane / 8) % 2)) * kSmemRowBytes + 16 * (lane / 16)); +#pragma unroll + for (int kd = 0; kd < kHeadDimQk / 32; ++kd) + { + uint32_t a0, a1, a2, a3; + ldmatrixX4(kAddr + 32 * kd, a0, a1, a2, a3); + mmaE4m3AsF16(acc[kd % 2], a0, a1, a2, a3, qh[kd]); + } + bool const valid0 = (validMask >> g) & 1U; + bool const valid1 = (validMask >> (g + 8)) & 1U; + s[0] = valid0 ? (acc[0][0] + acc[1][0]) * scaleSoftmaxLog2 : -INFINITY; + s[1] = valid0 ? (acc[0][1] + acc[1][1]) * scaleSoftmaxLog2 : -INFINITY; + s[2] = valid1 ? (acc[0][2] + acc[1][2]) * scaleSoftmaxLog2 : -INFINITY; + s[3] = valid1 ? (acc[0][3] + acc[1][3]) * scaleSoftmaxLog2 : -INFINITY; + } + + // The CTA's max per head column; this thread holds columns 2c and 2c + 1. + float max0 = fmaxf(s[0], s[2]); + float max1 = fmaxf(s[1], s[3]); +#pragma unroll + for (int offset = 4; offset < 32; offset *= 2) + { + max0 = fmaxf(max0, __shfl_xor_sync(0xffffffffU, max0, offset)); + max1 = fmaxf(max1, __shfl_xor_sync(0xffffffffU, max1, offset)); + } + if (g == 0) + { + sRedMax[warp * kMmaHeads + 2 * c] = max0; + sRedMax[warp * kMmaHeads + 2 * c + 1] = max1; + } + // Every warp observed its own gather barrier before its QK; this barrier orders all token rows before the P.V + // reads of the other warps. + __syncthreads(); + if (warp == 1) + { + // v_b_proj[h][16 r, 16 r + 16)[0, 512) for every head h: a weight, so it may be read at any point. It is + // issued once every token row has landed and arrives during QK / P.V. + if (lane == 0) + { + mbarrierArriveExpectTx(barVb, kVbBytes); + } + __syncwarp(); + if (lane < kHeads) + { + bulkCopyGlobalToShared(smemAddress(smem + kOffVb + lane * kVbHeadBytes), + static_cast(params.vbProj) + + (static_cast(lane) * kVHeadDim + split * kVPerCta) * kHeadDimV * 2, + kVbHeadBytes, barVb); + } + } + float ctaMax0 = -INFINITY; + float ctaMax1 = -INFINITY; +#pragma unroll + for (int w = 0; w < kWarps; ++w) + { + ctaMax0 = fmaxf(ctaMax0, sRedMax[w * kMmaHeads + 2 * c]); + ctaMax1 = fmaxf(ctaMax1, sRedMax[w * kMmaHeads + 2 * c + 1]); + } + // A CTA without valid tokens keeps P = 0. + float const offset0 = ctaMax0 == -INFINITY ? 0.f : ctaMax0; + float const offset1 = ctaMax1 == -INFINITY ? 0.f : ctaMax1; + + // P in E4M3 at the permuted k slots, and the per-head sums of the fp32 P. + { + float const p00 = exp2f(s[0] - offset0); + float const p01 = exp2f(s[1] - offset1); + float const p10 = exp2f(s[2] - offset0); + float const p11 = exp2f(s[3] - offset1); + int const slot0 = permutedSlot(tokBase + g); + int const slot1 = permutedSlot(tokBase + g + 8); + sP[(2 * c) * kPRowBytes + slot0] = toE4m3(p00 * kPScale); + sP[(2 * c + 1) * kPRowBytes + slot0] = toE4m3(p01 * kPScale); + sP[(2 * c) * kPRowBytes + slot1] = toE4m3(p10 * kPScale); + sP[(2 * c + 1) * kPRowBytes + slot1] = toE4m3(p11 * kPScale); + float sum0 = p00 + p10; + float sum1 = p01 + p11; +#pragma unroll + for (int offset = 4; offset < 32; offset *= 2) + { + sum0 += __shfl_xor_sync(0xffffffffU, sum0, offset); + sum1 += __shfl_xor_sync(0xffffffffU, sum1, offset); + } + if (g == 0) + { + sRedSum[warp * kMmaHeads + 2 * c] = sum0; + sRedSum[warp * kMmaHeads + 2 * c + 1] = sum1; + } + } + __syncthreads(); + + // This CTA's max and sum per head column go to every CTA of the cluster now; the partial output follows during + // P.V. The merged slice and the gathered output have buffers of their own, so no merge can overwrite memory that a + // warp of the cluster still reads in P.V. + clusterWait(); + if (tid < kSplits * kMmaHeads) + { + int const dst = tid / kMmaHeads; + int const h = tid % kMmaHeads; + float ctaMax = -INFINITY; + float ctaSum = 0.f; +#pragma unroll + for (int w = 0; w < kWarps; ++w) + { + ctaMax = fmaxf(ctaMax, sRedMax[w * kMmaHeads + h]); + ctaSum += sRedSum[w * kMmaHeads + h]; + } + storeRemote(mapToRank(smemAddress(recvStats + (split * kMmaHeads + h) * 2), dst), ctaMax, ctaSum, + mapToRank(barRecv, dst)); + } + + // O^T = V^T P^T: warp w owns the 16-column output blocks 4w .. 4w + 3, one block at a time over all k steps, each + // block sent to its merging CTA while the next one is computed. ldmatrix.trans of 8 tokens x 16 bytes gives + // (token 2c, columns 2g and 2g + 1) and (token 2c + 1, same columns); the byte permutes regroup them into a V^T row + // (one column, four tokens), so fragment row g holds column 2g and row g + 8 holds column 2g + 1. + { + uint32_t const pAddr = smemAddress(sP + (lane % 8) * kPRowBytes + 16 * ((lane / 8) % 2)); + uint32_t const vAddr = smemAddress(sKv + lane * kSmemRowBytes + 16 * (4 * warp)); + uint32_t pb[kTokensPerCta / 32][4]; +#pragma unroll + for (int kStep = 0; kStep < kTokensPerCta / 32; ++kStep) + { + uint32_t b0, b1; + ldmatrixX2(pAddr + 32 * kStep, b0, b1); + unpackE4m3B(pb[kStep], b0, b1); + } +#pragma unroll + for (int bi = 0; bi < 4; ++bi) + { + float ob[4] = {0.f, 0.f, 0.f, 0.f}; +#pragma unroll + for (int kStep = 0; kStep < kTokensPerCta / 32; ++kStep) + { + uint32_t r0, r1, r2, r3; + ldmatrixX4Trans(vAddr + 32 * kStep * kSmemRowBytes + 16 * bi, r0, r1, r2, r3); + uint32_t const a0 = __byte_perm(r0, r1, 0x6420); + uint32_t const a1 = __byte_perm(r0, r1, 0x7531); + uint32_t const a2 = __byte_perm(r2, r3, 0x6420); + uint32_t const a3 = __byte_perm(r2, r3, 0x7531); + mmaE4m3AsF16(ob, a0, a1, a2, a3, pb[kStep]); + } + // Columns [32 r, 32 r + 32) of the partial output go to CTA r. + int const col = 16 * (4 * warp + bi) + 2 * g; + int const dst = col / kDvPerCta; + int const colInSlice = col % kDvPerCta; + uint32_t const remoteBar = mapToRank(barRecv, dst); + storeRemote(mapToRank(smemAddress(recvO + (split * kMmaHeads + 2 * c) * kDvPerCta + colInSlice), dst), + ob[0], ob[2], remoteBar); + storeRemote(mapToRank(smemAddress(recvO + (split * kMmaHeads + 2 * c + 1) * kDvPerCta + colInSlice), dst), + ob[1], ob[3], remoteBar); + } + } + // The dependents (the o_proj GEMM) wait for the whole grid before reading out; launching them now lets them stage + // their weights during the merge and the v_b product. + cudaTriggerProgrammaticLaunchCompletion(); + + // Merge the kSplits partials of this CTA's output slice: thread (h, column) of the 8 x 32 slice. Warp h merges head + // column h; lane r < kSplits computes the weight of partial r once, and the loop takes the weights from the lanes + // in partial order. + mbarrierWaitCluster(barRecv, 0); + int const h = tid / kDvPerCta; + int const colInSlice = tid % kDvPerCta; + float value; + { + float const laneMax = lane < kSplits ? recvStats[(lane * kMmaHeads + h) * 2] : -INFINITY; + float mergedMax = laneMax; +#pragma unroll + for (int offset = 16; offset > 0; offset /= 2) + { + mergedMax = fmaxf(mergedMax, __shfl_xor_sync(0xffffffffU, mergedMax, offset)); + } + float const laneAlpha = laneMax == -INFINITY ? 0.f : exp2f(laneMax - mergedMax); + float mergedSum = 0.f; + float acc = 0.f; +#pragma unroll + for (int r = 0; r < kSplits; ++r) + { + float const alpha = __shfl_sync(0xffffffffU, laneAlpha, r); + mergedSum += alpha * recvStats[(r * kMmaHeads + h) * 2 + 1]; + acc += alpha * recvO[(r * kMmaHeads + h) * kDvPerCta + colInSlice]; + } + value = mergedSum > 0.f ? acc / (mergedSum * kPScale) * outputScale : 0.f; + } + + // The merged slice in BF16 (the attention output the v_b product reads) goes to every CTA of the cluster. The merge + // maps thread tid to head column tid / 32, so each warp holds and sends one head row. + reinterpret_cast<__nv_bfloat16*>(sSlice)[h * kDvPerCta + colInSlice] = __float2bfloat16_rn(value); + __syncwarp(); + { + int const chunk = lane % 4; + uint4 const packed = *reinterpret_cast(sSlice + (h * kDvPerCta + 8 * chunk) * 2); + uint32_t const localAddr = smemAddress(sGatheredO + (h * kHeadDimV + split * kDvPerCta + 8 * chunk) * 2); +#pragma unroll + for (int k = 0; k < 2; ++k) + { + uint32_t const dst = lane / 4 + 8 * k; + storeRemote16(mapToRank(localAddr, dst), packed, mapToRank(barGatheredO, dst)); + } + } + mbarrierWaitCluster(barGatheredO, 0); + mbarrierWait(barVb, 0); + if (h >= kHeads) + { + // A head column without a v_b slice (warp-uniform: the merge maps warp h to head column h). + return; + } + + // out[t][h * 256 + 16 r + v] = sum_d O[h][d] v_b[h][16 r + v][d] for v = 0 .. 15: warp h, one m16n8k16 row tile of + // the 16 v rows (1 KiB apart). Per pair of k steps a thread loads 16 bytes of its rows g and g + 8; rows with odd g + // take the other 32-column half of every 64-column block, so the eight rows of a load cover all banks. B column n + // follows the same parity as rows with g = n, so the dot product of row g is D[g][g % 2] (thread c = 0). + { + int const parity = g % 2; + uint8_t const* rowG = smem + kOffVb + h * kVbHeadBytes + g * (kHeadDimV * 2) + 16 * c; + uint8_t const* rowG8 = rowG + 8 * (kHeadDimV * 2); + uint8_t const* oRow = sGatheredO + h * (kHeadDimV * 2) + 16 * c; + float acc[4][4]; +#pragma unroll + for (int a = 0; a < 4; ++a) + { +#pragma unroll + for (int i = 0; i < 4; ++i) + { + acc[a][i] = 0.f; + } + } +#pragma unroll + for (int kPair = 0; kPair < kHeadDimV / 32; ++kPair) + { + int const block = (kPair ^ parity) * 64; + uint4 const ra = *reinterpret_cast(rowG + block); + uint4 const rb = *reinterpret_cast(rowG8 + block); + uint4 const ob = *reinterpret_cast(oRow + block); + mmaBf16(acc[(2 * kPair) % 4], ra.x, rb.x, ra.y, rb.y, ob.x, ob.y); + mmaBf16(acc[(2 * kPair + 1) % 4], ra.z, rb.z, ra.w, rb.w, ob.z, ob.w); + } + if (c == 0) + { + float sums[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + sums[i] = ((acc[0][i] + acc[1][i]) + acc[2][i]) + acc[3][i]; + } + __nv_bfloat16* outRow = static_cast<__nv_bfloat16*>(params.out) + + static_cast(tokenQ) * params.outStride + h * kVHeadDim + split * kVPerCta; + outRow[g] = __float2bfloat16_rn(parity ? sums[1] : sums[0]); + outRow[g + 8] = __float2bfloat16_rn(parity ? sums[3] : sums[2]); + } + } +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +template +void launchGlmSparseMlaFlashDecode(GlmSparseMlaFlashDecodeParams const& params, cudaStream_t stream) +{ + auto const kernel = glmSparseMlaFlashDecodeKernel; + static bool const configured = [kernel]() + { + TLLM_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes)); + TLLM_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1)); + return true; + }(); + (void) configured; + + cudaLaunchConfig_t config; + config.gridDim = dim3(kSplits, params.numTokensQ, 1); + config.blockDim = dim3(kThreads, 1, 1); + config.dynamicSmemBytes = kSmemBytes; + config.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kSplits; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.attrs = attrs; + config.numAttrs = 2; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, params)); +} + +} // namespace + +void invokeGlmSparseMlaFlashDecode(GlmSparseMlaFlashDecodeParams const& params, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(params.topK > 0 && params.topK <= kGlmSparseMlaMaxTopK, + "GLM sparse MLA flash decode supports TopK up to %d, got %d", kGlmSparseMlaMaxTopK, params.topK); + TLLM_CHECK_WITH_INFO(params.seqLenQ > 0 && params.indicesStride >= params.topK && params.numTokensQ >= 0 + && params.numTokensQ % params.seqLenQ == 0, + "Invalid GLM sparse MLA geometry"); + TLLM_CHECK_WITH_INFO(params.vbProj != nullptr && params.out != nullptr && params.outStride >= kHeads * kVHeadDim, + "GLM sparse MLA flash decode needs v_b_proj and an output of %d columns per token", kHeads * kVHeadDim); + if (params.numTokensQ == 0) + { + return; + } + bool const remap = params.remapBlockTable != nullptr; + if (remap) + { + TLLM_CHECK_WITH_INFO(params.remapTokensPerBlock > 0 && params.remapMaxBlocks > 0 + && params.remapBlockTableStride >= params.remapMaxBlocks, + "Invalid GLM sparse MLA index remap geometry"); + } + TLLM_CHECK_WITH_INFO(!params.gatherBeforeWait || remap, + "The GLM sparse MLA flash decode gathers before its wait only with the index remap"); + if (remap && params.gatherBeforeWait) + { + launchGlmSparseMlaFlashDecode(params, stream); + } + else if (remap) + { + launchGlmSparseMlaFlashDecode(params, stream); + } + else + { + launchGlmSparseMlaFlashDecode(params, stream); + } +} + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.h new file mode 100644 index 000000000000..7620932c477f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.h @@ -0,0 +1,105 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +// The GLM-5.2 decode shape at TP16 the flash-decode sparse MLA serves: 4 query heads per rank sharing one latent KV +// head, QK head dim 576 (512 latent + 64 rope), latent V dim 512, a TopK of at most 2048 selected tokens per query +// token, and the v_b BMM (v_head_dim 256) folded into the output. +constexpr int32_t kGlmSparseMlaNumHeads = 4; +constexpr int32_t kGlmSparseMlaHeadDimQk = 576; +constexpr int32_t kGlmSparseMlaHeadDimV = 512; +constexpr int32_t kGlmSparseMlaMaxTopK = 2048; +constexpr int32_t kGlmSparseMlaVHeadDim = 256; + +struct GlmSparseMlaFlashDecodeParams +{ + // [numTokensQ, 4, 576] E4M3. + void const* q; + // KV pool base: pool token row r holds 576 E4M3 values at byte offset r * 576. + void const* kvPool; + // [numTokensQ, indicesStride] selected tokens of each query token (pool rows, or positions with the remap); + // negative entries are masked. + int32_t const* indices; + // [numTokensQ / seqLenQ] KV length of each sequence, this step's tokens included. + int32_t const* seqLensKv; + // Device scalar that maps Q.K^T to the base-2 softmax domain. + float const* scaleSoftmaxLog2; + // Device scalar that scales the attention output. + float const* outputScale; + // [4, 256, 512] BF16 v_b_proj. + void const* vbProj; + // out[t][h * 256 + v] = sum_d O[t][h][d] * vbProj[h][v][d] (BF16, outStride elements between query tokens), where + // O is the attention output rounded to BF16. + void* out; + int64_t outStride; + int32_t numTokensQ; + // Query tokens per sequence: query token t of a sequence attends to the first + // min(seqLensKv - (seqLenQ - 1 - t), topK) of its indices. + int32_t seqLenQ; + int32_t indicesStride; + int32_t topK; + // Optional in-kernel remap of the indices, with the arithmetic of convertReqIndexToGlobal: with remapBlockTable + // set, indices hold token positions in the query token's sequence, and position p selects pool token row + // remapBlockTable[s * remapBlockTableStride + p / remapTokensPerBlock] * remapStrideFactor + // + p % remapTokensPerBlock + remapLayerOffset, where s = tokenQ / seqLenQ is the sequence (the row of seqLensKv + // and of the block table). A negative position, a block at or beyond remapMaxBlocks or a negative table entry + // masks the entry. + int32_t const* remapBlockTable; + int64_t remapBlockTableStride; + int32_t remapMaxBlocks; + int32_t remapTokensPerBlock; + int32_t remapStrideFactor; + int32_t remapLayerOffset; + // With the remap, true when the indices were written by an earlier layer's top-k (a layer that shares it): the + // kernel then reads the indices, the block table and the KV length before its grid-dependency wait and gathers the + // rows of the positions before this step's tokens (position < seqLensKv - seqLenQ, appended by earlier steps) + // there too; the rows of this step's tokens, written by the KV append it depends on, after the wait. + bool gatherBeforeWait; +}; + +// Generation-only static-token sparse MLA with E4M3 Q and KV and the v_b BMM folded in, as flash decoding over the +// gathered keys. Every query token gets one cluster of 16 CTAs, and CTA s owns indices [128 s, 128 s + 128): +// - each warp gathers 16 token rows with one bulk copy per row onto its own mbarrier; +// - S = Q K^T and P V run as FP16 mma.sync on E4M3 values unpacked exactly to FP16, with fp32 accumulation; the MMAs +// run 8 head columns, the 4 query heads in columns 0-3 and zero Q rows in 4-7, whose outputs are not stored; +// P = exp2(S * scale - max) is stored as E4M3 scaled by 448, as the TRTLLM-Gen generation kernels feed their second +// GEMM, and the query fragments are unpacked once per warp; +// - every CTA sends output slice r of its unnormalized partial output (one 16-column block at a time, during P.V), its +// max and its sum straight into the shared memory of CTA r of the cluster (st.async), and CTA r merges the 16 +// partials of its 32 output columns; +// - CTA r also owns v in [16 r, 16 r + 16) of every head: it all-gathers the row's merged BF16 output (st.async), +// multiplies it with its v_b_proj slice (64 KiB, bulk-copied once every token row has landed) on mma.sync BF16, and +// writes those 64 values of out. +// Requires the SM100 family, numTokensQ * 16 CTAs of 170 KiB shared memory and clusters of 16 CTAs. PDL: the kernel +// reads global memory only after its grid-dependency wait (except the gatherBeforeWait reads above) and lets its +// dependents launch once its partial outputs are sent. +void invokeGlmSparseMlaFlashDecode(GlmSparseMlaFlashDecodeParams const& params, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.cu new file mode 100644 index 000000000000..4fd734cab81d --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.cu @@ -0,0 +1,122 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ +namespace +{ + +constexpr int kAcceptBlockSize = 128; + +// One CTA; a thread per request. The forced-acceptance counter is read and advanced once per call. +template +__global__ void __launch_bounds__(kAcceptBlockSize) verifyHeadAcceptKernel(VerifyHeadAcceptParams params) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + // The target tokens are written by the preceding kernel. + cudaGridDependencySynchronize(); +#endif + __shared__ int64_t counter; + bool const randomForce = params.forceMode == VerifyHeadForceMode::kRandom; + if (randomForce && threadIdx.x == 0) + { + counter = *params.forceCounter + 1; + *params.forceCounter = counter; + } + __syncthreads(); + + auto const* target = static_cast(params.targetTokens); + int const batchSize = params.numContexts + params.numGens; + int const rowLen = params.draftLen + 1; + for (int req = static_cast(threadIdx.x); req < batchSize; req += kAcceptBlockSize) + { + int64_t const rowOffset = static_cast(req) * rowLen; + if (req < params.numContexts) + { + params.numAcceptedTokens[req] = 1; + if (params.acceptedTokens != nullptr) + { + params.acceptedTokens[rowOffset] = static_cast(target[req]); + } + continue; + } + int const gen = req - params.numContexts; + TargetT const* genTarget = target + params.numContexts + static_cast(gen) * rowLen; + int32_t const* genDraft = params.draftTokens + static_cast(gen) * params.draftLen; + // The target token plus the leading draft tokens that equal the target tokens. + int accepted = 1; + while (accepted <= params.draftLen + && static_cast(genDraft[accepted - 1]) == static_cast(genTarget[accepted - 1])) + { + ++accepted; + } + if (params.forceMode == VerifyHeadForceMode::kFixed) + { + accepted = params.forceBase; + } + else if (randomForce) + { + int64_t const poolIdx + = (counter * params.counterStride + static_cast(gen) * params.slotStride) & params.poolMask; + int const extra = params.forcePool[poolIdx] < params.forceFrac ? 1 : 0; + accepted = min(params.forceBase + extra, params.forceMax); + } + params.numAcceptedTokens[req] = accepted; + if (params.acceptedTokens != nullptr) + { + for (int j = 0; j < rowLen; ++j) + { + params.acceptedTokens[rowOffset + j] = static_cast(genTarget[j]); + } + } + } +} + +} // namespace + +void invokeVerifyHeadAccept(VerifyHeadAcceptParams const& params) +{ + TLLM_CHECK_WITH_INFO(params.forceMode != VerifyHeadForceMode::kRandom + || (params.forcePool != nullptr && params.forceCounter != nullptr), + "verifyHeadAccept: the random forced count needs the pool and the counter"); + if (params.numContexts + params.numGens == 0) + { + return; + } + if (params.targetTokensInt64) + { + tensorrt_llm::common::launchWithPdlWhenEnabled("verifyHeadAccept", verifyHeadAcceptKernel, dim3(1), + dim3(kAcceptBlockSize), 0, params.stream, params); + } + else + { + tensorrt_llm::common::launchWithPdlWhenEnabled("verifyHeadAccept", verifyHeadAcceptKernel, dim3(1), + dim3(kAcceptBlockSize), 0, params.stream, params); + } +} + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.h new file mode 100644 index 000000000000..52cd9212e7c6 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.h @@ -0,0 +1,72 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +//! How the accepted counts of the generation requests are forced (TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS). +enum class VerifyHeadForceMode : int32_t +{ + kNone = 0, //!< The strict acceptance count + kFixed = 1, //!< forceBase + kRandom = 2, //!< min(forceBase + (pool[(counter * counterStride + gen * slotStride) & poolMask] < forceFrac), + //!< forceMax), counter advanced by one per call +}; + +//! \brief Inputs and outputs of the fused strict + forced acceptance of one-model speculative decoding. +struct VerifyHeadAcceptParams +{ + void const* targetTokens; //!< [numContexts + numGens * (draftLen + 1)], int32, or int64 when targetTokensInt64 + bool targetTokensInt64; + int32_t const* draftTokens; //!< [numGens, draftLen] + int32_t* numAcceptedTokens; //!< Output [numContexts + numGens] + int32_t* acceptedTokens; //!< Output [numContexts + numGens, draftLen + 1], or nullptr: not written + int numContexts; + int numGens; + int draftLen; + VerifyHeadForceMode forceMode; + int forceBase; + int forceMax; + float forceFrac; + float const* forcePool; //!< kRandom: fixed uniform pool + int64_t* forceCounter; //!< kRandom: call counter + int64_t counterStride; + int64_t slotStride; + int64_t poolMask; + cudaStream_t stream; +}; + +//! \brief One kernel for SpecWorkerBase._sample_and_accept_draft_tokens_base after sampling. +//! +//! Context rows accept their sampled token (count 1, accepted row [target, ...]). A generation row accepts its target +//! token plus the leading draft tokens equal to the target tokens (the cumprod count of the strict acceptance), and +//! then the forced count replaces it as forceMode says. Integer results equal the ATen ops bit for bit. Under PDL +//! the kernel waits for its producer before it reads memory and lets its dependents launch at its end. +void invokeVerifyHeadAccept(VerifyHeadAcceptParams const& params); + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/Gemm2K128.cu b/cpp/tensorrt_llm/kernels/moe/trtllmGen/Gemm2K128.cu new file mode 100644 index 000000000000..57b00b763d83 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/Gemm2K128.cu @@ -0,0 +1,666 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// FC2 of the NVFP4 MoE at intermediate size 128 (K = 128) and tile N 8 (see Gemm2K128.h for the contract). +// +// Work. A unit is (token tile t, 128-row tile of the output features); the units are split evenly over a persistent +// grid of one CTA per SM, each CTA owning a contiguous run of units (one or two token tiles). A CTA holds every unit of +// its run at once: the weight tiles (8 KB of E2M1, SWIZZLE_64B, by TMA in two-tile boxes) and their 128x4 block scales +// (one bulk copy per token tile) are requested before the grid-dependency wait, the activations of its token tiles +// (8 rows x 64 B) and their 8x4 scales after it. +// +// MMA. tcgen05.mma kind::mxf4nvf4 block16, M = 128 (weight rows), N = 8 (token rows), two K = 64 steps per unit, issued +// by the elected lane of the converged MMA warp, one commit per unit; the weight scales go into TMEM by tcgen05.cp +// ahead of the MMAs (a separate warp, as their bulk copy lands), the token scales after the wait. The TRTLLM-Gen cubin +// computes the same two K steps and two more on the zero padding of its 256-wide K tile; adding +0.0f to the +// accumulator gives its result, bitwise. +// +// Accumulators. Unit k of a token tile's run writes its 8 TMEM columns at base + stride * k, stride = the tile's valid +// rows rounded up to 2, 4 or 8 (accumulators are 2-column aligned): the valid rows of all units of the tile lie in +// stride * units consecutive columns, and the tail columns a unit shares with the next one are padding rows, written +// in issue order. The epilogue reads them in 16-column loads, each once its last unit is complete. +// +// Warps: 0-3 epilogue (TMEM lanes 32 w .. 32 w + 31 = weight rows of each tile), 4 weight loads, 5 TMEM allocation, +// activations and MMA, 6 weight-scale copies. + +#include "Gemm2K128.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaDriverWrapper.h" +#include "tensorrt_llm/common/cudaUtils.h" + +#include +#include + +namespace moe::dev::gemm2K128 +{ +namespace +{ + +constexpr int kTileRows = 128; +constexpr int kK = 128; +constexpr int kTokens = 8; +constexpr int kTileBytesA = kTileRows * kK / 2; // 8,192: 128 rows x 64 B, SWIZZLE_64B +constexpr int kTileBytesSf = kTileRows * kK / 16; // 1,024: two 512 B blocks of the 128x4 layout +constexpr int kActBytes = kTokens * kK / 2; // 512: one SWIZZLE_64B atom +constexpr int kSfbBytes = 512; // 32 rows x 16 B: the tcgen05.cp source of a token tile's scales +constexpr int kMaxUnits = 15; // units per CTA +constexpr int kMaxSlots = 2; // token tiles per CTA (a run of kMaxUnits units spans at most two) +constexpr int kMaxTokenTiles = 64; +constexpr int kThreads = 224; +constexpr uint32_t kTmemCols = 512; +constexpr uint32_t kTmemAcc = 0; // packed accumulators, at most 256 columns +constexpr uint32_t kTmemSfa = 256; // kMaxUnits x 8 columns: the weight scales of every unit +constexpr uint32_t kTmemSfb = 376; // kMaxSlots x 4 columns: per K step a column of token scales and a zero column +// Instruction descriptor, kind::mxf4nvf4: A/B E2M1, UE4M3 scales, K-major, M = 128, N = 8. +constexpr uint32_t kIdesc = (1u << 7) | (1u << 10) | ((8u >> 3) << 17) | ((128u >> 4) << 24); +// K-major SWIZZLE_64B operand descriptor without its start address: 8-row x 64 B atoms 512 B apart, version 1. +constexpr uint64_t kDescSw64 = (1ull << 16) | (static_cast(512 >> 4) << 32) | (1ull << 46) | (4ull << 61); +// tcgen05.cp 32x128b source descriptor without its start address: 8-row x 16 B core matrices 128 B apart. +constexpr uint64_t kDescSf = (static_cast(128 >> 4) << 32) | (1ull << 46); +constexpr int kSmemBytes = kMaxUnits * (kTileBytesA + kTileBytesSf) + kMaxSlots * (kActBytes + kSfbBytes) + 1024; + +struct Params +{ + CUtensorMap tmW2; // weights [experts * hidden rows][64 B], box {64, 256}, SWIZZLE_64B + CUtensorMap tmW1; // the same tensor, box {64, 128} + uint8_t const* wSf; + uint8_t const* act; + uint8_t const* actSf; + float const* scale; + int32_t const* numNonExitingCtas; + int32_t const* ctaIdxXyToBatchIdx; + int32_t const* ctaIdxXyToMnLimit; + __nv_bfloat16* out; + int32_t hidden; + int32_t maxCtas; +}; + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1100 + +__device__ __forceinline__ uint32_t smemU32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ void mbarInit(uint64_t* bar, uint32_t count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" ::"r"(smemU32(bar)), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarExpectTx(uint64_t* bar, uint32_t bytes) +{ + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(smemU32(bar)), "r"(bytes) : "memory"); +} + +__device__ __forceinline__ void mbarWait(uint64_t* bar, uint32_t parity) +{ + uint32_t ok = 0; + while (!ok) + { + asm volatile( + "{\n" + ".reg .pred p;\n" + "mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n" + "selp.u32 %0, 1, 0, p;\n" + "}\n" + : "=r"(ok) + : "r"(smemU32(bar)), "r"(parity) + : "memory"); + } +} + +__device__ __forceinline__ void tmaLoad2d(void* dst, void const* map, int x, int y, uint64_t* bar) +{ + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1, {%2, %3}], [%4];" ::"r"( + smemU32(dst)), + "l"(map), "r"(x), "r"(y), "r"(smemU32(bar)) + : "memory"); +} + +__device__ __forceinline__ void bulkLoad(void* dst, void const* src, uint32_t bytes, uint64_t* bar) +{ + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"(smemU32(dst)), + "l"(src), "r"(bytes), "r"(smemU32(bar)) + : "memory"); +} + +__device__ __forceinline__ int ldRelaxed(int32_t const* p) +{ + int v; + asm volatile("ld.relaxed.gpu.global.b32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ void namedBarSync(int id, int threads) +{ + asm volatile("bar.sync %0, %1;" ::"r"(id), "r"(threads) : "memory"); +} + +__device__ __forceinline__ void namedBarArrive(int id, int threads) +{ + asm volatile("bar.arrive %0, %1;" ::"r"(id), "r"(threads) : "memory"); +} + +__device__ __forceinline__ void fenceProxyAsyncSmem() +{ + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + +__device__ __forceinline__ void tcFenceBefore() +{ + asm volatile("tcgen05.fence::before_thread_sync;" ::: "memory"); +} + +__device__ __forceinline__ void tcFenceAfter() +{ + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); +} + +__device__ __forceinline__ void tmemAlloc(uint32_t* dst, uint32_t cols) +{ + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"(smemU32(dst)), "r"(cols) + : "memory"); +} + +__device__ __forceinline__ void tmemRelinquish() +{ + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;" ::: "memory"); +} + +__device__ __forceinline__ void tmemDealloc(uint32_t taddr, uint32_t cols) +{ + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(taddr), "r"(cols) : "memory"); +} + +// The following three are executed by a whole (converged) warp; its elected lane issues the instruction. +__device__ __forceinline__ void utccpElect(uint32_t taddr, uint64_t sdesc) +{ + asm volatile( + "{\n" + ".reg .pred e;\n" + "elect.sync _|e, 0xffffffff;\n" + "@e tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;\n" + "}\n" ::"r"(taddr), + "l"(sdesc) + : "memory"); +} + +__device__ __forceinline__ void mmaElect(uint32_t d, uint64_t a, uint64_t b, uint32_t sfa, uint32_t sfb, uint32_t acc) +{ + asm volatile( + "{\n" + ".reg .pred e, p;\n" + "elect.sync _|e, 0xffffffff;\n" + "setp.ne.b32 p, %6, 0;\n" + "@e tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%4], [%5], p;\n" + "}\n" ::"r"(d), + "l"(a), "l"(b), "r"(kIdesc), "r"(sfa), "r"(sfb), "r"(acc) + : "memory"); +} + +__device__ __forceinline__ void commitElect(uint64_t* bar) +{ + asm volatile( + "{\n" + ".reg .pred e;\n" + "elect.sync _|e, 0xffffffff;\n" + "@e tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];\n" + "}\n" ::"r"(smemU32(bar)) + : "memory"); +} + +// Columns [taddr, taddr + 16) of the warp's TMEM lane quarter. +__device__ __forceinline__ void tmemLd16(uint32_t taddr, float (&v)[16]) +{ + uint32_t r[16]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32 {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, " + "%15}, [%16];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), "=r"(r[8]), + "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]) + : "r"(taddr) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" + : "+r"(r[0]), "+r"(r[1]), "+r"(r[2]), "+r"(r[3]), "+r"(r[4]), "+r"(r[5]), "+r"(r[6]), "+r"(r[7]), + "+r"(r[8]), "+r"(r[9]), "+r"(r[10]), "+r"(r[11]), "+r"(r[12]), "+r"(r[13]), "+r"(r[14]), "+r"(r[15]) + : + : "memory"); +#pragma unroll + for (int i = 0; i < 16; ++i) + { + v[i] = __uint_as_float(r[i]); + } +} + +// Stores the valid rows of the units [kBase, kBase + 16 / kStride) of a token tile from 16 packed accumulator columns +// (column kStride * k + r = row r of unit kBase + k). rowPtr: row 0 of the tile at this thread's feature of the run's +// first unit. +template +__device__ __forceinline__ void storeChunk( + float const (&v)[16], __nv_bfloat16* rowPtr, size_t rowStride, int kBase, int runLen, int nv, float scale) +{ + constexpr int kUnits = 16 / kStride; +#pragma unroll + for (int k = 0; k < kUnits; ++k) + { + if (kBase + k < runLen) + { + __nv_bfloat16* o = rowPtr + (kBase + k) * kTileRows; +#pragma unroll + for (int r = 0; r < kStride; ++r) + { + if (r < nv) + { + o[r * rowStride] = __float2bfloat16_rn(__fadd_rn(v[kStride * k + r], 0.0f) * scale); + } + } + } + } +} + +__device__ __forceinline__ int unitStart(int c, int units, int grid) +{ + return static_cast(static_cast(c) * units / grid); +} + +#endif + +__global__ void __launch_bounds__(kThreads, 1) gemm2K128Kernel(const __grid_constant__ Params p) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1100 + extern __shared__ __align__(1024) uint8_t smemRaw[]; + __shared__ __align__(8) uint64_t fullBar[kMaxUnits]; + __shared__ __align__(8) uint64_t sfaBar[kMaxSlots]; + __shared__ __align__(8) uint64_t sfaReady[kMaxUnits]; + __shared__ __align__(8) uint64_t accFull[kMaxUnits]; + __shared__ uint32_t sTmem; + __shared__ int sRange[2]; + __shared__ int sExpert[kMaxSlots]; + __shared__ int sNv[kMaxSlots]; + __shared__ int sStride[kMaxSlots]; + __shared__ int sCol[kMaxSlots]; + + uint8_t* smem = reinterpret_cast((reinterpret_cast(smemRaw) + 1023) & ~uintptr_t(1023)); + uint8_t* smemA = smem; + uint8_t* smemSf = smemA + kMaxUnits * kTileBytesA; + uint8_t* smemAct = smemSf + kMaxUnits * kTileBytesSf; + uint8_t* smemSfb = smemAct + kMaxSlots * kActBytes; + + int const warp = threadIdx.x >> 5; + int const lane = threadIdx.x & 31; + int const tilesPerSlot = p.hidden / kTileRows; + + if (warp == 4 && lane == 0) + { + for (int i = 0; i < kMaxUnits; ++i) + { + mbarInit(&fullBar[i], 1); + mbarInit(&sfaReady[i], 1); + mbarInit(&accFull[i], 1); + } + for (int j = 0; j < kMaxSlots; ++j) + { + mbarInit(&sfaBar[j], 1); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + asm volatile("prefetch.tensormap [%0];" ::"l"(&p.tmW2) : "memory"); + asm volatile("prefetch.tensormap [%0];" ::"l"(&p.tmW1) : "memory"); + } + if (warp == 5) + { + tmemAlloc(&sTmem, kTmemCols); + tmemRelinquish(); + } + tcFenceBefore(); + __syncthreads(); + tcFenceAfter(); + uint32_t const tmem = sTmem; + + if (warp == 4) + { + // The routing outputs are complete when this kernel starts: its primary, the FC1 cubin, triggers its dependents + // only after its own grid-dependency wait on the routing kernel. One round trip for the tile count and the + // experts and valid rows of the token tiles; the weights and scales are constants. + int const numTiles = ldRelaxed(p.numNonExitingCtas); + int const e0 = lane < p.maxCtas ? ldRelaxed(p.ctaIdxXyToBatchIdx + lane) : 0; + int const e1 = lane + 32 < p.maxCtas ? ldRelaxed(p.ctaIdxXyToBatchIdx + lane + 32) : 0; + int const m0 = lane < p.maxCtas ? ldRelaxed(p.ctaIdxXyToMnLimit + lane) : 0; + int const m1 = lane + 32 < p.maxCtas ? ldRelaxed(p.ctaIdxXyToMnLimit + lane + 32) : 0; + int const units = numTiles * tilesPerSlot; + int const a = unitStart(blockIdx.x, units, gridDim.x); + int const b = unitStart(blockIdx.x + 1, units, gridDim.x); + // Publish the run (its experts, valid rows and accumulator columns) before issuing the loads. + if (b > a) + { + int const tA = a / tilesPerSlot; + int const tB = (b - 1) / tilesPerSlot; + int col = 0; + for (int t = tA; t <= tB; ++t) + { + int const eLo = __shfl_sync(0xffffffffu, e0, t & 31); + int const eHi = __shfl_sync(0xffffffffu, e1, t & 31); + int const mLo = __shfl_sync(0xffffffffu, m0, t & 31); + int const mHi = __shfl_sync(0xffffffffu, m1, t & 31); + int const nv = (t < 32 ? mLo : mHi) - kTokens * t; + int const stride = nv <= 2 ? 2 : (nv <= 4 ? 4 : 8); + int const runLen = min(b, (t + 1) * tilesPerSlot) - max(a, t * tilesPerSlot); + if (lane == 0) + { + sExpert[t - tA] = t < 32 ? eLo : eHi; + sNv[t - tA] = nv; + sStride[t - tA] = stride; + sCol[t - tA] = col; + } + // The next token tile starts at a 32-column boundary (aligned TMEM loads). + col = (col + stride * runLen + 31) & ~31; + } + } + if (lane == 0) + { + sRange[0] = a; + sRange[1] = b; + } + __syncwarp(); + namedBarArrive(1, kThreads); + // Per token tile of the run: all its block scales in one bulk copy (first, for the scale-copy warp), then its + // weight tiles in two-tile boxes (a one-tile box for an odd last one). + if (lane == 0) + { + int i = 0; + int slot = 0; + while (a + i < b) + { + int const u0 = a + i; + int const t = u0 / tilesPerSlot; + int const nt0 = u0 - t * tilesPerSlot; + int const n = min(b - u0, tilesPerSlot - nt0); + mbarExpectTx(&sfaBar[slot], static_cast(n * kTileBytesSf)); + bulkLoad(smemSf + i * kTileBytesSf, + p.wSf + (static_cast(sExpert[slot]) * tilesPerSlot + nt0) * kTileBytesSf, + static_cast(n * kTileBytesSf), &sfaBar[slot]); + i += n; + ++slot; + } + i = 0; + slot = 0; + while (a + i < b) + { + int const u0 = a + i; + int const t = u0 / tilesPerSlot; + int const nt0 = u0 - t * tilesPerSlot; + int const n = min(b - u0, tilesPerSlot - nt0); + int const row0 = sExpert[slot] * p.hidden + nt0 * kTileRows; + for (int k = 0; k < n; k += 2) + { + bool const two = k + 1 < n; + mbarExpectTx(&fullBar[i + k], two ? 2 * kTileBytesA : kTileBytesA); + tmaLoad2d(smemA + (i + k) * kTileBytesA, two ? &p.tmW2 : &p.tmW1, 0, row0 + k * kTileRows, + &fullBar[i + k]); + } + i += n; + ++slot; + } + } + } + else + { + namedBarSync(1, kThreads); + } + int const a = sRange[0]; + int const b = sRange[1]; + int const n = b - a; + int const t0 = n > 0 ? a / tilesPerSlot : 0; + int const nt0 = a - t0 * tilesPerSlot; + int const numSlots = n > 0 ? (b - 1) / tilesPerSlot - t0 + 1 : 0; + + if (warp == 6) + { + // The weight scales of every unit into TMEM as their bulk copy lands (before the grid-dependency wait returns, + // for a CTA that runs while the previous kernel does); the MMA warp waits on sfaReady[unit]. + int u = 0; + int nt = nt0; + for (int j = 0; j < numSlots; ++j) + { + int const runLen = min(n - u, tilesPerSlot - nt); + mbarWait(&sfaBar[j], 0); + tcFenceAfter(); + for (int k = 0; k < runLen; ++k) + { + uint32_t const src = smemU32(smemSf + (u + k) * kTileBytesSf); + uint32_t const dst = tmem + kTmemSfa + 8 * (u + k); + utccpElect(dst, kDescSf | ((src & 0x3ffffu) >> 4)); + utccpElect(dst + 4, kDescSf | (((src + 512) & 0x3ffffu) >> 4)); + commitElect(&sfaReady[u + k]); + } + u += runLen; + nt = 0; + } + } + + // Everything below reads FC1's output or writes this kernel's output (the weight-load and scale-copy warps read + // neither). The dependents may launch once every CTA is past its wait. + if (warp < 4 || warp == 5) + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + if (threadIdx.x == 0) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + + if (warp < 4) + { + // Epilogue, per token tile of the run, in 16-column chunks, each once its last unit is complete. + uint32_t const tq = tmem + (static_cast(32 * warp) << 16); + // The weight rows are shuffled for the 128-row epilogue tile: row 32 w + l of a tile holds its feature + // 32 w + 4 (l & 7) + l / 8. + int const feature = 32 * warp + (lane & 7) * 4 + (lane >> 3); + int runStart = 0; + for (int j = 0; j < numSlots; ++j) + { + int const t = t0 + j; + int const ntFirst = j == 0 ? nt0 : 0; + int const runLen = min(n - runStart, tilesPerSlot - ntFirst); + int const nv = sNv[j]; + int const stride = sStride[j]; + float const scale = __ldg(p.scale + sExpert[j]); + int const unitsPerChunk = 16 / stride; + __nv_bfloat16* const rowBase + = p.out + static_cast(t) * kTokens * p.hidden + ntFirst * kTileRows + feature; + for (int kBase = 0; kBase < runLen; kBase += unitsPerChunk) + { + mbarWait(&accFull[runStart + min(kBase + unitsPerChunk, runLen) - 1], 0); + tcFenceAfter(); + float v[16]; + tmemLd16(tq + kTmemAcc + sCol[j] + stride * kBase, v); + if (stride == 2) + { + storeChunk<2>(v, rowBase, p.hidden, kBase, runLen, nv, scale); + } + else if (stride == 4) + { + storeChunk<4>(v, rowBase, p.hidden, kBase, runLen, nv, scale); + } + else + { + storeChunk<8>(v, rowBase, p.hidden, kBase, runLen, nv, scale); + } + } + runStart += runLen; + } + } + else if (warp == 5 && n > 0) + { + // Activations of the run's token tiles (8 rows of K = 128 E2M1, SWIZZLE_64B) and their 8x4 block scales, all + // loaded in one round trip. The scales go through shared memory (rows 0-7 of a 32-row block: per K step a + // column of 4 scales and a zero column) into every TMEM lane quarter. + int const r8 = lane >> 2; + int const cnk = lane & 3; + uint4 actV[kMaxSlots]; + uint32_t sf0[kMaxSlots]; + uint32_t sf1[kMaxSlots]; +#pragma unroll + for (int j = 0; j < kMaxSlots; ++j) + { + actV[j] = make_uint4(0, 0, 0, 0); + sf0[j] = 0; + sf1[j] = 0; + if (j < numSlots) + { + actV[j] = __ldcg(reinterpret_cast( + p.act + (static_cast(t0 + j) * kTokens + r8) * (kK / 2) + cnk * 16)); + if (lane < kTokens) + { + sf0[j] = __ldcg(reinterpret_cast(p.actSf) + 2 * (t0 + j) * kTokens + lane); + sf1[j] = __ldcg(reinterpret_cast(p.actSf) + (2 * (t0 + j) + 1) * kTokens + lane); + } + } + } +#pragma unroll + for (int j = 0; j < kMaxSlots; ++j) + { + if (j < numSlots) + { + *reinterpret_cast(smemAct + j * kActBytes + r8 * 64 + ((cnk ^ ((r8 >> 1) & 3)) << 4)) = actV[j]; + *reinterpret_cast(smemSfb + j * kSfbBytes + lane * 16) = make_uint4(sf0[j], 0u, sf1[j], 0u); + } + } + fenceProxyAsyncSmem(); + __syncwarp(); + for (int j = 0; j < numSlots; ++j) + { + utccpElect(tmem + kTmemSfb + 4 * j, kDescSf | ((smemU32(smemSfb + j * kSfbBytes) & 0x3ffffu) >> 4)); + } + + int slot = 0; + int nt = nt0; + int runIdx = 0; // index of the unit in its token tile's run + uint32_t aBase = smemU32(smemA); + uint32_t bBase = smemU32(smemAct); + uint32_t sfb = tmem + kTmemSfb; + uint32_t d = tmem + kTmemAcc + sCol[0]; + uint32_t stride = sStride[0]; + for (int i = 0; i < n; ++i) + { + if (nt == tilesPerSlot) + { + ++slot; + nt = 0; + runIdx = 0; + bBase += kActBytes; + sfb += 4; + d = tmem + kTmemAcc + sCol[slot]; + stride = sStride[slot]; + } + if ((runIdx & 1) == 0) + { + // The first unit of a two-tile box. + mbarWait(&fullBar[i], 0); + } + mbarWait(&sfaReady[i], 0); + tcFenceAfter(); + uint32_t const sfa = tmem + kTmemSfa + 8 * i; + uint64_t const da = kDescSw64 | ((aBase & 0x3ffffu) >> 4); + uint64_t const db = kDescSw64 | ((bBase & 0x3ffffu) >> 4); + mmaElect(d, da, db, sfa, sfb, 0u); + mmaElect(d, da + 2, db + 2, sfa + 4, sfb + 2, 1u); + // A commit per unit, although the epilogue waits only for the last unit of each chunk: with a commit per + // chunk instead, the issue of a unit takes about 1.4x longer. + commitElect(&accFull[i]); + ++nt; + ++runIdx; + aBase += kTileBytesA; + d += stride; + } + } + + tcFenceBefore(); + __syncthreads(); + if (warp == 5) + { + tcFenceAfter(); + tmemDealloc(tmem, kTmemCols); + } +#endif +} + +void encodeWeightMap(CUtensorMap* map, void const* weights, uint64_t rows, uint32_t boxRows) +{ + cuuint64_t const dims[2] = {64, rows}; + cuuint64_t const strides[1] = {64}; + cuuint32_t const box[2] = {64, boxRows}; + cuuint32_t const elementStrides[2] = {1, 1}; + TLLM_CU_CHECK(tensorrt_llm::common::CUDADriverWrapper::getInstance()->cuTensorMapEncodeTiled(map, + CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(weights), dims, strides, box, elementStrides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_64B, CU_TENSOR_MAP_L2_PROMOTION_L2_128B, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); +} + +} // namespace + +bool isSupported(int32_t hiddenSize, int32_t maxNumCtas) +{ + int const sm = tensorrt_llm::common::getSMVersion(); + if (sm < 100 || sm >= 110 || hiddenSize % kTileRows != 0 || hiddenSize / kTileRows < kMaxUnits || maxNumCtas < 1 + || maxNumCtas > kMaxTokenTiles) + { + return false; + } + int const grid = tensorrt_llm::common::getMultiProcessorCount(); + int const maxUnits = maxNumCtas * (hiddenSize / kTileRows); + return (maxUnits + grid - 1) / grid <= kMaxUnits; +} + +void run(Data const& data, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(isSupported(data.hiddenSize, data.maxNumCtas), "gemm2K128: unsupported shape."); + TLLM_CUDA_CHECK(cudaFuncSetAttribute(gemm2K128Kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes)); + Params p{}; + uint64_t const rows = static_cast(data.numLocalExperts) * data.hiddenSize; + encodeWeightMap(&p.tmW2, data.weights, rows, 2 * kTileRows); + encodeWeightMap(&p.tmW1, data.weights, rows, kTileRows); + p.wSf = static_cast(data.weightsScale); + p.act = static_cast(data.act); + p.actSf = static_cast(data.actScale); + p.scale = data.outputScale; + p.numNonExitingCtas = data.numNonExitingCtas; + p.ctaIdxXyToBatchIdx = data.ctaIdxXyToBatchIdx; + p.ctaIdxXyToMnLimit = data.ctaIdxXyToMnLimit; + p.out = static_cast<__nv_bfloat16*>(data.output); + p.hidden = data.hiddenSize; + p.maxCtas = data.maxNumCtas; + + cudaLaunchConfig_t config{}; + config.gridDim = dim3(tensorrt_llm::common::getMultiProcessorCount()); + config.blockDim = dim3(kThreads); + config.dynamicSmemBytes = kSmemBytes; + config.stream = stream; + cudaLaunchAttribute attributes[2]{}; + attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attributes[0].val.programmaticStreamSerializationAllowed = data.usePdl ? 1 : 0; + attributes[1].id = cudaLaunchAttributePreferredSharedMemoryCarveout; + attributes[1].val.sharedMemCarveout = static_cast(data.sharedMemCarveout); + config.attrs = attributes; + config.numAttrs = data.sharedMemCarveout >= 0 ? 2 : 1; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, gemm2K128Kernel, p)); +} + +} // namespace moe::dev::gemm2K128 diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/Gemm2K128.h b/cpp/tensorrt_llm/kernels/moe/trtllmGen/Gemm2K128.h new file mode 100644 index 000000000000..72d99550e18f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/Gemm2K128.h @@ -0,0 +1,60 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +namespace moe::dev::gemm2K128 +{ + +// FC2 of the NVFP4 MoE at intermediate size 128 and tile N 8, with the inputs and output of the TRTLLM-Gen Gemm2 +// cubins: +// output[8 t + r][f] = bf16(outputScale[e] * sum_k weights[e][f][k] * act[8 t + r][k]), e = ctaIdxXyToBatchIdx[t], +// for the routing's token tiles t < numNonExitingCtas and their valid rows r < ctaIdxXyToMnLimit[t] - 8 t. +struct Data +{ + // [numLocalExperts * hiddenSize][64 B] E2M1, the Gemm2 weight layout (rows shuffled for the 128-row epilogue tile). + void const* weights; + // E4M3 block scales of the weights, 128x4 layout: [numLocalExperts][hiddenSize / 128][2][512 B]. + void const* weightsScale; + // FC1 output: [rows][64 B] E2M1 and its E4M3 block scales in the 8x4 layout ([rows / 8][2][8][4]). + void const* act; + void const* actScale; + float const* outputScale; // [numLocalExperts] + int32_t const* numNonExitingCtas; + int32_t const* ctaIdxXyToBatchIdx; + int32_t const* ctaIdxXyToMnLimit; + void* output; // [rows][hiddenSize] BF16 + int32_t hiddenSize; + int32_t numLocalExperts; + int32_t maxNumCtas; // entries of ctaIdxXyToBatchIdx and ctaIdxXyToMnLimit + bool usePdl; + // The launch's preferred shared-memory carveout in percent (cudaLaunchAttributePreferredSharedMemoryCarveout; + // cudaSharedmemCarveoutMaxShared = 100), or -1: no preference, the driver chooses. + int32_t sharedMemCarveout = -1; +}; + +// Whether run() covers the shape on this device: an SM100-family GPU, hiddenSize a multiple of 128 and at least 1,920, +// at most 64 token tiles, and at most 15 units (token tile x 128 output features) per SM. +bool isSupported(int32_t hiddenSize, int32_t maxNumCtas); + +void run(Data const& data, cudaStream_t stream); + +} // namespace moe::dev::gemm2K128 diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingFusedRouter.cu b/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingFusedRouter.cu new file mode 100644 index 000000000000..7e447cdceb4a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingFusedRouter.cu @@ -0,0 +1,516 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Router GEMV + DeepSeek-style routing (sigmoid + bias, top-8 of 256 experts, scaled sum normalization) in one kernel +// for 1-8 tokens, with the outputs of routingIndicesSmallTokenKernel. +// +// Logits. Every logit is the value dsv3RouterGemm.cu's router_gemm_kernel computes (bit for bit). There, thread t +// (0..127) of the CTA of expert e accumulates acc = fma(a, w, acc) over k = ki * 1024 + 8 t + j (ki = 0..5, j = 0..7), +// in that order; warp s (the "slice" of threads 32 s .. 32 s + 31) adds its lanes with the xor butterfly 16, 8, 4, 2, 1 +// (lane 0's value), and the four slice sums are added to 0.0f in slice order. Here a GEMV CTA covers one slice of 8 +// experts: it stages the slice's activations of every token in shared memory (24 KB at 8 tokens), runs the same +// per-thread chains (fma.rn.f32.bf16 computes fma(float(a), float(w), acc) without conversion instructions), adds the +// 32 positions with the same pairs as the butterfly and publishes the slice sums. 128 CTAs cover 256 experts x 4 +// slices, and each CTA stages a quarter of the activations instead of all of them. +// +// Handoff. The slice sums go to one of two exchange buffers in the workspace (by the parity of the epoch word) whose +// idle value is -0.0. A sum of -0.0 is published as +0.0, which gives the same logit, since the four sums are added to +// 0.0f. The routing CTA polls its token rows until no -0.0 is left, so no fence or counter is needed. It resets the +// other buffer, consumed by the previous launch, while it waits, and advances the epoch at its end. Launches that share +// a workspace must therefore run in stream order. +// +// Routing. Phase A of routingIndicesSmallTokenKernel (one warp per token: sigmoid + bias, the packed-key top-K, the +// scaled sum normalization; the same operations in the same order). The metadata then follows from the set of +// selected local experts after a single barrier: with no more tokens than tile rows every selected local expert owns +// exactly one CTA tile, so its tile index is its rank among the selected local experts and its rows start at +// rank * tile, which is what the small-token kernel's scans compute. + +#include "RoutingCustomPolicy.cuh" +#include "RoutingKernel.h" + +#include + +namespace moe::dev::routing +{ +namespace routingCustom +{ +namespace +{ + +using FusedExpertSelect = TopKExpertSelect; +using FusedKernelParams = KernelParams; + +constexpr int kNumExperts = 256; +constexpr int kTopK = 8; +constexpr int kMaxTokens = 8; +constexpr int kHiddenDim = 6144; +constexpr int kEltsPerVec = 8; // bf16 per 16-byte load, as in the router GEMV +constexpr int kKIterElts = kEltsPerVec * 128; // elements per k-iteration of the router GEMV +constexpr int kKIters = kHiddenDim / kKIterElts; // 6 +constexpr int kSlices = 4; // warps per expert in the router GEMV +constexpr int kExpertsPerCta = 8; +constexpr int kNumGemvCtas = kNumExperts / kExpertsPerCta * kSlices; // 128 +constexpr int kNumThreads = 256; +constexpr int kSliceVecs = kKIters * WarpSize; // 16-byte vectors of one token's slice +constexpr int kActBytes = kMaxTokens * kSliceVecs * 16; // 24 KB +constexpr int kPartBytes = WarpSize * kExpertsPerCta * kMaxTokens * 4; // 8 KB +constexpr int kSmemBytes = kActBytes + kPartBytes; +constexpr int kExchangeFloats = kMaxTokens * kNumExperts * kSlices; +constexpr int kKIdxRowBytes = 8; +constexpr uint32_t kNegZeroBits = 0x80000000u; + +static_assert(kNumExperts == kNumThreads, "the routing CTA runs one thread per expert"); +static_assert(2 * kExchangeFloats + 4 == FusedRouterWorkspaceWords, "workspace layout"); + +struct FusedRouterParams +{ + FusedKernelParams routing; + __nv_bfloat16 const* hiddenStates; + __nv_bfloat16 const* routerWeight; + int32_t* workspace; +}; + +__device__ __forceinline__ uint32_t smemAddr(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ void gridDependencyWait() +{ + // The memory clobber keeps the compiler from moving memory accesses across the wait. + asm volatile("griddepcontrol.wait;" ::: "memory"); +} + +// Non-coherent 16-byte load of router weights. +__device__ __forceinline__ uint4 loadWeight16(void const* ptr) +{ + uint4 v; + asm volatile("ld.global.nc.v4.u32 {%0, %1, %2, %3}, [%4];" : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(ptr)); + return v; +} + +__device__ __forceinline__ uint32_t loadVolatile(uint32_t const* ptr) +{ + uint32_t v; + asm volatile("ld.volatile.global.u32 %0, [%1];" : "=r"(v) : "l"(ptr) : "memory"); + return v; +} + +__device__ __forceinline__ float4 loadVolatile4(float const* ptr) +{ + float4 v; + asm volatile("ld.volatile.global.v4.f32 {%0, %1, %2, %3}, [%4];" + : "=f"(v.x), "=f"(v.y), "=f"(v.z), "=f"(v.w) + : "l"(ptr) + : "memory"); + return v; +} + +// A slice sum for the routing CTA, which polls it with volatile loads. +__device__ __forceinline__ void storeRelaxed(float* ptr, float v) +{ + asm volatile("st.relaxed.gpu.global.f32 [%0], %1;" ::"l"(ptr), "f"(v) : "memory"); +} + +__device__ __forceinline__ bool isNegZero(float v) +{ + return __float_as_uint(v) == kNegZeroBits; +} + +// acc = fma(a[j], w[j], acc) for j = 0..7, in that order, on 8 packed bf16 pairs. +__device__ __forceinline__ void fmaBf16x8(float& acc, uint4 const& a, uint4 const& w) +{ + uint32_t const aw[4] = {a.x, a.y, a.z, a.w}; + uint32_t const ww[4] = {w.x, w.y, w.z, w.w}; +#pragma unroll + for (int q = 0; q < 4; q++) + { + asm("{\n" + ".reg .b16 al, ah, bl, bh;\n" + "mov.b32 {al, ah}, %1;\n" + "mov.b32 {bl, bh}, %2;\n" + "fma.rn.f32.bf16 %0, al, bl, %0;\n" + "fma.rn.f32.bf16 %0, ah, bh, %0;\n" + "}\n" + : "+f"(acc) + : "r"(aw[q]), "r"(ww[q])); + } +} + +// The slice GEMV of one CTA: slice r of experts 8 g .. 8 g + 7. Thread mapping: warp w covers positions 8 (w & 3) .. +// + 7 of the slice and experts 4 (w >> 2) .. + 3, so the 4 expert groups of a warp read the same 128 bytes of +// activations. Returns through emit(eLocal, m, sliceSum) once per (expert, token). +template +__device__ __forceinline__ void sliceGemv(uint8_t* smem, __nv_bfloat16 const* hiddenStates, uint4 const (&w)[kKIters], + int r, int t, int eLocal, Emit const& emit) +{ + uint4* act = reinterpret_cast(smem); + float* part = reinterpret_cast(smem + kActBytes); + + // Stage the slice: token m, k-iteration ki, positions 32 r .. 32 r + 31. + constexpr int kVecs = kNumTokens * kSliceVecs; +#pragma unroll + for (int j = 0; j < (kVecs + kNumThreads - 1) / kNumThreads; ++j) + { + int const c = threadIdx.x + kNumThreads * j; + if (c < kVecs) + { + int const m = c / kSliceVecs; + int const rem = c - m * kSliceVecs; + __nv_bfloat16 const* src = hiddenStates + m * kHiddenDim + (rem / WarpSize) * kKIterElts + + (WarpSize * r + rem % WarpSize) * kEltsPerVec; + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(smemAddr(&act[c])), "l"(src) : "memory"); + } + } + asm volatile("cp.async.commit_group;\ncp.async.wait_group 0;" ::: "memory"); + __syncthreads(); + + // The per-thread chains of router thread 32 r + t for one expert and every token. + float acc[kMaxTokens] = {}; +#pragma unroll + for (int ki = 0; ki < kKIters; ki++) + { +#pragma unroll + for (int m = 0; m < kNumTokens; m++) + { + fmaBf16x8(acc[m], act[(m * kKIters + ki) * WarpSize + t], w[ki]); + } + } + + // The butterfly of the slice: 4 threads per (expert, token) fold the positions {sub, sub + 4, ..., sub + 28} + // through levels 16, 8 and 4, two shuffles add levels 2 and 1. + float4* row = reinterpret_cast(part + (t * kExpertsPerCta + eLocal) * kMaxTokens); + row[0] = make_float4(acc[0], acc[1], acc[2], acc[3]); + row[1] = make_float4(acc[4], acc[5], acc[6], acc[7]); + __syncthreads(); + int const pair = threadIdx.x / 4; + int const sub = threadIdx.x % 4; + int const eL = pair / kMaxTokens; + int const m = pair % kMaxTokens; + auto v = [&](int tt) { return part[(tt * kExpertsPerCta + eL) * kMaxTokens + m]; }; + float const s1a = v(sub) + v(sub + 16); + float const s1b = v(sub + 4) + v(sub + 20); + float const s1c = v(sub + 8) + v(sub + 24); + float const s1d = v(sub + 12) + v(sub + 28); + float sum = (s1a + s1c) + (s1b + s1d); + sum += __shfl_xor_sync(0xffffffff, sum, 2); + sum += __shfl_xor_sync(0xffffffff, sum, 1); + if (sub == 0 && m < kNumTokens) + { + emit(eL, m, sum); + } +} + +template +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) +__global__ void __launch_bounds__(kNumThreads, 1) routingFusedRouterKernel(FusedRouterParams const __grid_constant__ p) +{ + extern __shared__ __align__(128) uint8_t smemSlice[]; + __shared__ __align__(8) uint8_t smemKIdx[kNumExperts * kKIdxRowBytes]; + __shared__ int16_t smemTopKExpert[kMaxTokens * kTopK]; + __shared__ float smemBias[kNumExperts]; + __shared__ __align__(16) uint32_t smemSelected[kNumExperts / 32]; + + auto const& params = p.routing; + float* const buffers = reinterpret_cast(p.workspace); + uint32_t* const epochPtr = reinterpret_cast(p.workspace + 2 * kExchangeFloats); + + if (blockIdx.x < kNumGemvCtas) + { + int const r = blockIdx.x % kSlices; + int const g = blockIdx.x / kSlices; + int const warp = threadIdx.x / WarpSize; + int const lane = threadIdx.x % WarpSize; + int const t = 8 * (warp % 4) + lane % 8; + int const eLocal = 4 * (warp / 4) + lane / 8; + // The weights are not written by the primary grid: load them before the wait. A shared-memory store of their + // XOR ahead of a barrier makes the loads complete before the wait instead of being scheduled after it. + uint4 w[kKIters]; + __nv_bfloat16 const* wRow + = p.routerWeight + (kExpertsPerCta * g + eLocal) * kHiddenDim + (WarpSize * r + t) * kEltsPerVec; + uint32_t loaded = 0; +#pragma unroll + for (int ki = 0; ki < kKIters; ki++) + { + w[ki] = loadWeight16(wRow + ki * kKIterElts); + loaded ^= w[ki].x ^ w[ki].y ^ w[ki].z ^ w[ki].w; + } + reinterpret_cast(smemSlice + kActBytes)[threadIdx.x] = loaded; + __syncthreads(); + gridDependencyWait(); + // The previous launch on this workspace precedes the primary grid in the stream, so it has completed. + float* const exchange = buffers + (loadVolatile(epochPtr) & 1u) * kExchangeFloats; + sliceGemv(smemSlice, p.hiddenStates, w, r, t, eLocal, + [&](int eL, int m, float sliceSum) + { + storeRelaxed(&exchange[(m * kNumExperts + kExpertsPerCta * g + eL) * kSlices + r], + isNegZero(sliceSum) ? 0.f : sliceSum); + }); + // Exit without a trigger: the dependents launch once the routing CTA triggers. + return; + } + + // Routing CTA. Independent of the primary grid: clear this expert's row of the byte map, stage its bias. + *reinterpret_cast(&smemKIdx[threadIdx.x * kKIdxRowBytes]) = uint2{0xFFFFFFFFu, 0xFFFFFFFFu}; + auto const& pre = params.mExpertSelectParams.mPreprocessParams; + smemBias[threadIdx.x] = loadScalar(pre.ptrRoutingBias, threadIdx.x, pre.dtypeBias); + if (threadIdx.x < kNumExperts / 32) + { + smemSelected[threadIdx.x] = 0u; + } + __syncthreads(); + gridDependencyWait(); + uint32_t const epoch = loadVolatile(epochPtr); + float* const exchange = buffers + (epoch & 1u) * kExchangeFloats; + { + // The buffer of the previous launch, which its routing CTA has consumed: idle again for the next launch. + float4* other = reinterpret_cast(buffers + ((epoch + 1u) & 1u) * kExchangeFloats); + for (int i = threadIdx.x; i < kExchangeFloats / 4; i += kNumThreads) + { + other[i] = make_float4(-0.f, -0.f, -0.f, -0.f); + } + } + + // ----- Phase A: one warp per token, as in routingIndicesSmallTokenKernel ----- + int32_t const warpIdx = __shfl_sync(0xffffffff, threadIdx.x / WarpSize, 0); + int32_t const laneIdx = cutlass::arch::LaneId(); + auto block = cg::this_thread_block(); + auto warp = cg::tiled_partition(block); + if (warpIdx < kNumTokens) + { + constexpr int VecSize = kNumExperts / WarpSize; + float4 slices[VecSize]; + bool ready; + do + { + ready = true; +#pragma unroll + for (int i = 0; i < VecSize; ++i) + { + slices[i] = loadVolatile4(exchange + (warpIdx * kNumExperts + i * WarpSize + laneIdx) * kSlices); + ready = ready && !isNegZero(slices[i].x) && !isNegZero(slices[i].y) && !isNegZero(slices[i].z) + && !isNegZero(slices[i].w); + } + } while (!ready); + + float score[VecSize]; + int32_t idx[VecSize]; +#pragma unroll + for (int i = 0; i < VecSize; ++i) + { + idx[i] = i * WarpSize + laneIdx; + float logit = 0.0f; + logit += slices[i].x; + logit += slices[i].y; + logit += slices[i].z; + logit += slices[i].w; + score[i] = sigmoid_accurate(logit) + smemBias[idx[i]]; + } + float topScore[kTopK]; + int32_t topExpert[kTopK]; + topk::reduceTopK(warp, topScore, topExpert, score, idx, float{-INFINITY}, kTopK); + + // Lane k's top-K entry, the value the small-token kernel reads as warpTopKScore[laneIdx]. + float myScore = 0.f; + int32_t myExpert = 0; +#pragma unroll + for (int k = 0; k < kTopK; ++k) + { + myScore = laneIdx == k ? topScore[k] : myScore; + myExpert = laneIdx == k ? topExpert[k] : myExpert; + } + bool const isTop = laneIdx < kTopK; + float const sigmoidScore = isTop ? (myScore - smemBias[myExpert]) : 0.f; + // cg::reduce's butterfly (xor 16, 8, 4, 2, 1) adds +0.0 to lanes 0-7 in its first two levels, which leaves + // them unchanged (no sigmoidScore is -0.0); levels 4, 2 and 1 then pair lanes 0-7 as here. + float sum = sigmoidScore; + sum += __shfl_xor_sync(0xffffffff, sum, 4); + sum += __shfl_xor_sync(0xffffffff, sum, 2); + sum += __shfl_xor_sync(0xffffffff, sum, 1); + if (isTop) + { + auto const& post = params.mExpertSelectParams.mPostprocessParams; + float const weight = sigmoidScore * post.routeScale / (sum + post.sumEpsilon); + smemKIdx[myExpert * kKIdxRowBytes + warpIdx] = static_cast(laneIdx); + smemTopKExpert[warpIdx * kTopK + laneIdx] = static_cast(myExpert); + atomicOr(&smemSelected[myExpert / 32], 1u << (myExpert % 32)); + params.mPtrTopKWeights[warpIdx * kTopK + laneIdx] = __nv_bfloat16{weight}; + } + } + __syncthreads(); + + // ----- Metadata: tile index = rank among the selected local experts ----- + int32_t const localBegin = params.mLocalExpertsStartIdx; + int32_t const localEnd = params.mLocalExpertsStartIdx + params.mNumLocalExperts; + uint4 const lo = *reinterpret_cast(&smemSelected[0]); + uint4 const hi = *reinterpret_cast(&smemSelected[4]); + uint32_t selected[8] = {lo.x, lo.y, lo.z, lo.w, hi.x, hi.y, hi.z, hi.w}; + if (localBegin != 0 || localEnd < kNumExperts) + { +#pragma unroll + for (int i = 0; i < 8; ++i) + { + int32_t const b0 = max(localBegin - 32 * i, 0); + int32_t const b1 = min(localEnd - 32 * i, 32); + uint32_t const upTo = b1 >= 32 ? 0xFFFFFFFFu : ((1u << max(b1, 0)) - 1u); + uint32_t const below = b0 >= 32 ? 0xFFFFFFFFu : ((1u << b0) - 1u); + selected[i] &= b1 <= b0 ? 0u : (upTo & ~below); + } + } + int32_t numSelected = 0; +#pragma unroll + for (int i = 0; i < 8; ++i) + { + numSelected += __popc(selected[i]); + } + // (rank, is selected and local) of an expert, from masked popcounts of every word: indexing the words by the expert + // would place them in local memory. + auto rankOf = [&](int32_t e, bool& isSelected) + { + int32_t const word = e >> 5; + uint32_t const bit = 1u << (e & 31); + int32_t rank = 0; + int32_t hit = 0; +#pragma unroll + for (int i = 0; i < 8; ++i) + { + rank += __popc(selected[i] & (i < word ? 0xFFFFFFFFu : (i == word ? bit - 1u : 0u))); + hit += __popc(selected[i] & (i == word ? bit : 0u)); + } + isSelected = hit != 0; + return rank; + }; + auto tileRows = [&](int32_t n) + { return params.mIsPow2 ? mulLog2(n, params.mPaddingLog2) : mulTileN(n, params.mTileTokensDim); }; + + { + int32_t const expert = threadIdx.x; + bool isSelected; + int32_t const rank = rankOf(expert, isSelected); + if (isSelected) + { + uint2 const kRow = *reinterpret_cast(&smemKIdx[expert * kKIdxRowBytes]); + int32_t const count = __popc(__vsetne4(kRow.x, 0xFFFFFFFFu)) + __popc(__vsetne4(kRow.y, 0xFFFFFFFFu)); + params.mPtrCtaIdxXyToBatchIdx[rank] = expert - localBegin; + params.mPtrCtaIdxXyToMnLimit[rank] = min(tileRows(rank + 1), tileRows(rank) + count); + } + } + if (threadIdx.x == 0) + { + params.mPtrPermutedIdxSize[0] = tileRows(numSelected); + params.mPtrNumNonExitingCtas[0] = numSelected; + } + + // Permutation, one (token, k) pair per thread. offset-within-expert = number of earlier tokens that also selected + // the expert (its byte-map row below byte t), the small-token kernel's order. + if (static_cast(threadIdx.x) < kNumTokens * kTopK) + { + int32_t const expandedIdx = threadIdx.x; + int32_t const tokenIdx = expandedIdx / kTopK; + int32_t const pairExpert = smemTopKExpert[expandedIdx]; + bool pairIsLocal; + int32_t const rank = rankOf(pairExpert, pairIsLocal); + uint2 const row = *reinterpret_cast(&smemKIdx[pairExpert * kKIdxRowBytes]); + uint64_t const present = (static_cast(__vsetne4(row.y, 0xFFFFFFFFu)) << 32) + | static_cast(__vsetne4(row.x, 0xFFFFFFFFu)); + uint64_t const earlier = tokenIdx == 0 ? 0ull : (present & ((1ull << (8 * tokenIdx)) - 1)); + int32_t const permutedIdx = pairIsLocal ? tileRows(rank) + __popcll(earlier) : int32_t{-1}; + if (params.mPtrExpandedIdxToPermutedIdx != nullptr) + { + params.mPtrExpandedIdxToPermutedIdx[expandedIdx] = permutedIdx; + } + if (params.mPtrPermutedIdxToExpandedIdx != nullptr && pairIsLocal) + { + params.mPtrPermutedIdxToExpandedIdx[permutedIdx] = expandedIdx; + } + if (params.mPtrPermutedIdxToTokenIdx != nullptr && pairIsLocal) + { + params.mPtrPermutedIdxToTokenIdx[permutedIdx] = tokenIdx; + } + } + if (threadIdx.x == 0) + { + *epochPtr = epoch + 1u; + } + // After all routing outputs are written. + cudaTriggerProgrammaticLaunchCompletion(); +} +#else +__global__ void __launch_bounds__(kNumThreads, 1) + routingFusedRouterKernel(FusedRouterParams const __grid_constant__ /* p */) +{ + assert(false && "routingFusedRouterKernel is only supported on SM100+ architectures"); +} +#endif // if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) + +template +void launchFusedRouter(FusedRouterParams const& params, bool usePdl, cudaStream_t stream) +{ + cudaLaunchConfig_t config{}; + config.gridDim = kNumGemvCtas + 1; + config.blockDim = kNumThreads; + config.dynamicSmemBytes = kSmemBytes; + config.stream = stream; + cudaLaunchAttribute attributes[1] = {}; + attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attributes[0].val.programmaticStreamSerializationAllowed = int(usePdl); + config.attrs = attributes; + config.numAttrs = 1; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, routingFusedRouterKernel, params)); +} + +} // namespace + +bool supportsFusedRouter(Data const& data, int32_t hiddenDim) +{ + return hiddenDim == kHiddenDim && data.mNumExperts == kNumExperts && data.mTopK == kTopK && data.mNumTokens >= 1 + && data.mNumTokens <= kMaxTokens && data.mNumTokens <= data.mTileTokensDim && data.mLocalExpertsStrideLog2 == 0 + && data.mPreprocessType == RoutingPreprocessType::SigmoidBias + && data.mPostprocessType == RoutingPostprocessType::ScaledSumNormalize + && data.mDtypeOutput == tg::Dtype::Bfloat16 && data.mPtrTopKWeights != nullptr + && data.mPtrRoutingBias != nullptr && data.mPtrTopKIds == nullptr; +} + +void runFusedRouter(Data const& data, void const* hiddenStates, void const* routerWeight, int32_t hiddenDim, + int32_t* workspace, void* stream) +{ + TLLM_CHECK_WITH_INFO(supportsFusedRouter(data, hiddenDim), + "The fused router + routing kernel supports 1-8 tokens (at most the tile size), 256 experts, top-8, hidden " + "size %d and SigmoidBias + ScaledSumNormalize routing; got %d tokens, tile %d, %d experts, top-%d, hidden %d", + kHiddenDim, data.mNumTokens, data.mTileTokensDim, data.mNumExperts, data.mTopK, hiddenDim); + FusedRouterParams params; + params.routing = FusedKernelParams::setKernelParams(data); + params.hiddenStates = static_cast<__nv_bfloat16 const*>(hiddenStates); + params.routerWeight = static_cast<__nv_bfloat16 const*>(routerWeight); + params.workspace = workspace; + auto const s = static_cast(stream); + switch (data.mNumTokens) + { + case 1: launchFusedRouter<1>(params, data.mUsePdl, s); break; + case 2: launchFusedRouter<2>(params, data.mUsePdl, s); break; + case 3: launchFusedRouter<3>(params, data.mUsePdl, s); break; + case 4: launchFusedRouter<4>(params, data.mUsePdl, s); break; + case 5: launchFusedRouter<5>(params, data.mUsePdl, s); break; + case 6: launchFusedRouter<6>(params, data.mUsePdl, s); break; + case 7: launchFusedRouter<7>(params, data.mUsePdl, s); break; + default: launchFusedRouter<8>(params, data.mUsePdl, s); break; + } +} + +} // namespace routingCustom +} // namespace moe::dev::routing diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingKernel.h b/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingKernel.h index 3f238844e7fc..668aeb8b9ec2 100644 --- a/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingKernel.h +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingKernel.h @@ -384,6 +384,21 @@ void launchBlockKernel(Data const& data, uint32_t numThreadsHist, void* stream); void launchCoopBlockKernel(Data const& data, uint32_t numThreadsHist, void* stream); void launchSmallTokenBlockKernel(Data const& data, uint32_t numThreadsHist, void* stream); +// Workspace of runFusedRouter, in 32-bit words: two exchange buffers of [8 tokens][256 experts][4 slices] fp32 whose +// idle value is -0.0 (bits 0x80000000), then the epoch (0 before the first launch), padded to 16 bytes. +static constexpr int32_t FusedRouterWorkspaceWords = 2 * 8 * 256 * 4 + 4; + +// Whether runFusedRouter covers this routing configuration and hidden size: 1-8 tokens and at most mTileTokensDim, +// 256 experts, top-8, unit local-expert stride, SigmoidBias + ScaledSumNormalize, bf16 expert weights. +bool supportsFusedRouter(Data const& data, int32_t hiddenDim); + +// The router GEMV of dsv3_router_gemm_op (bf16 hiddenStates [mNumTokens, hiddenDim] x routerWeight [mNumExperts, +// hiddenDim], the same fp32 logits, which are not stored) and the routing of routingIndicesSmallTokenKernel in one +// kernel, with the same routing outputs; mPtrScores is not read. workspace: FusedRouterWorkspaceWords words, owned by +// the caller and initialized once; launches that share a workspace must run in stream order. +void runFusedRouter(Data const& data, void const* hiddenStates, void const* routerWeight, int32_t hiddenDim, + int32_t* workspace, void* stream); + } // namespace routingCustom //////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.cu b/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.cu index 01ec7600681c..fa2c2d9b6327 100644 --- a/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.cu +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.cu @@ -15,6 +15,7 @@ */ #include "DevKernel.h" +#include "Gemm2K128.h" #include "routing/RoutingKernel.h" #include "runner.h" #include "tensorrt_llm/common/config.h" @@ -411,6 +412,44 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3 serializeMoeRoutingMethodType(routingMethodType).c_str(), (int) routingMethodType); } } + +void Runner::runFusedRouter(void const* routerHiddenStates, void const* routerWeight, int32_t hiddenDim, + int32_t* routerWorkspace, void* routingBias, int32_t numTokens, int32_t numExperts, int32_t topK, + int32_t localExpertOffset, int32_t localNumExperts, float routedScalingFactor, int32_t* permutedIdxSize, + int32_t* expandedIdxToPermutedIdx, int32_t* permutedIdxToExpandedIdx, int32_t* permutedIdxToTokenIdx, + void* expertWeights, int32_t* ctaIdxXyToBatchIdx, int32_t* ctaIdxXyToMnLimit, int32_t* numNonExitingCtas, + cudaStream_t stream, btg::Dtype dtypeRoutingBias) +{ + // The routing configuration of run()'s DeepSeekV3 branch without groups. + moe::dev::routing::routingCustom::Data routingData; + routingData.mDtypeOutput = btg::Dtype::Bfloat16; + routingData.mDtypeInput = btg::Dtype::Fp32; + routingData.mUsePdl = tensorrt_llm::common::getEnvEnablePDL(); + routingData.mPreprocessType = moe::dev::routing::RoutingPreprocessType::SigmoidBias; + routingData.mPostprocessType = moe::dev::routing::RoutingPostprocessType::ScaledSumNormalize; + routingData.mPtrRoutingBias = routingBias; + routingData.mDtypeBias = dtypeRoutingBias; + routingData.mRouteScale = routedScalingFactor; + routingData.mPtrPermutedIdxSize = permutedIdxSize; + routingData.mPtrExpandedIdxToPermutedIdx = expandedIdxToPermutedIdx; + routingData.mPtrPermutedIdxToExpandedIdx = permutedIdxToExpandedIdx; + routingData.mPtrPermutedIdxToTokenIdx = permutedIdxToTokenIdx; + routingData.mPtrTopKWeights = expertWeights; + routingData.mPtrCtaIdxXyToBatchIdx = ctaIdxXyToBatchIdx; + routingData.mPtrCtaIdxXyToMnLimit = ctaIdxXyToMnLimit; + routingData.mPtrNumNonExitingCtas = numNonExitingCtas; + routingData.mNumTokens = numTokens; + routingData.mNumExperts = numExperts; + routingData.mTopK = topK; + routingData.mPaddingLog2 = computeLog2(mTileTokensDim); + routingData.mTileTokensDim = mTileTokensDim; + routingData.mLocalExpertsStartIdx = localExpertOffset; + routingData.mLocalExpertsStrideLog2 = 0; + routingData.mNumLocalExperts = localNumExperts; + + moe::dev::routing::routingCustom::runFusedRouter( + routingData, routerHiddenStates, routerWeight, hiddenDim, routerWorkspace, stream); +} } // namespace Routing namespace PermuteGemm1 @@ -675,6 +714,10 @@ Runner::Runner(btg::Dtype dtypeAct, btg::Dtype dtypeWeights, bool useDeepSeekFp8 , mGemm2(Gemm2::Runner(dtypeAct, dtypeWeights, btg::Dtype::Bfloat16, useDeepSeekFp8, tileTokensDim, useFineGrained)) , mActType(actType) , mTileTokensDim(tileTokensDim) + , mDtypeAct(dtypeAct) + , mDtypeWeights(dtypeWeights) + , mUseDeepSeekFp8(useDeepSeekFp8) + , mUseFineGrained(useFineGrained) { auto const& gemm1PassingIndices = mPermuteGemm1.getPassingConfigIndices(); auto const& gemm2PassingIndices = mGemm2.getPassingConfigIndices(); @@ -840,6 +883,11 @@ void Runner::run( auto const& config = mPassingConfigs[configIndex]; + if (args.hidden_states_ready != nullptr) + { + // gemm1 is the first kernel that reads hidden_states; routing, launched before this point, does not wait. + TLLM_CUDA_CHECK(cudaStreamWaitEvent(stream, args.hidden_states_ready, 0)); + } mPermuteGemm1.run(args.hidden_states, hidden_states_scale_linear, args.gemm1_weights, args.gemm1_weights_scale, workspace.expert_weights, args.output1_scales_scalar, args.output1_scales_gate_scalar, args.gemm1_bias, args.gemm1_alpha, args.gemm1_beta, args.gemm1_clamp_limit, workspace.gemm1_output, workspace.gemm1_output_scale, @@ -863,13 +911,50 @@ void Runner::run( } // Run gemm2 - mGemm2.run(gemm2_input, gemm2_input_scale, args.gemm2_weights, args.gemm2_weights_scale, args.output2_scales_scalar, - args.gemm2_bias, workspace.gemm2_output, workspace.gemm2_output_scale, totalExpertsPerToken, - args.output_hidden_size.value_or(args.hidden_size), args.intermediate_size, totalLocalExperts, args.num_tokens, - workspace.num_non_exiting_ctas, workspace.total_num_padded_tokens, workspace.cta_idx_xy_to_batch_idx, - workspace.cta_idx_xy_to_mn_limit, workspace.bmm2_workspace, device, stream, config.gemm2Config, - args.valid_hidden_size.value_or(args.hidden_size), - args.valid_intermediate_size.value_or(args.intermediate_size)); + int32_t const gemm2HiddenSize = args.output_hidden_size.value_or(args.hidden_size); + int32_t const gemm2MaxNumCtas + = Routing::getMaxNumCgasInBatchDim(args.num_tokens, totalExpertsPerToken, totalLocalExperts, mTileTokensDim); + // gemm2K128 covers the FC2 of the NVFP4 MoE at intermediate size 128 and tile N 8 (for example GLM-5 decode at MoE + // TP16): the same output rows as the Gemm2 cubin, bitwise (see Gemm2K128.cu). + if (args.prefer_gemm2_k128 && mDtypeAct == btg::Dtype::E2m1 && mDtypeWeights == btg::Dtype::E2m1 && !mUseDeepSeekFp8 + && !mUseFineGrained && mTileTokensDim == 8 && args.num_fused_shared_experts == 0 + && args.intermediate_size == 128 + && args.valid_intermediate_size.value_or(args.intermediate_size) == args.intermediate_size + && args.valid_hidden_size.value_or(args.hidden_size) == gemm2HiddenSize && args.gemm2_bias == nullptr + && args.output2_scales_scalar != nullptr && moe::dev::gemm2K128::isSupported(gemm2HiddenSize, gemm2MaxNumCtas)) + { + moe::dev::gemm2K128::Data gemm2Data{}; + gemm2Data.weights = args.gemm2_weights; + gemm2Data.weightsScale = args.gemm2_weights_scale; + gemm2Data.act = gemm2_input; + gemm2Data.actScale = gemm2_input_scale; + gemm2Data.outputScale = args.output2_scales_scalar; + gemm2Data.numNonExitingCtas = workspace.num_non_exiting_ctas; + gemm2Data.ctaIdxXyToBatchIdx = workspace.cta_idx_xy_to_batch_idx; + gemm2Data.ctaIdxXyToMnLimit = workspace.cta_idx_xy_to_mn_limit; + gemm2Data.output = workspace.gemm2_output; + gemm2Data.hiddenSize = gemm2HiddenSize; + gemm2Data.numLocalExperts = totalLocalExperts; + gemm2Data.maxNumCtas = gemm2MaxNumCtas; + gemm2Data.usePdl = tensorrt_llm::common::getEnvEnablePDL(); + // The maximum shared-memory carveout preferred. Without a preference the driver can run the 141 KB CTAs in the + // 164 KB configuration. The finalize kernel that gemm2K128 lets launch early (for example an all-reduce hop + // with the finalize fused) then shares those SMs and keeps them in that configuration until it ends, so the + // next layer's large-shared-memory CTAs that land there start only after it. A scheduling hint only: the + // outputs are unchanged. + gemm2Data.sharedMemCarveout = static_cast(cudaSharedmemCarveoutMaxShared); + moe::dev::gemm2K128::run(gemm2Data, stream); + } + else + { + mGemm2.run(gemm2_input, gemm2_input_scale, args.gemm2_weights, args.gemm2_weights_scale, + args.output2_scales_scalar, args.gemm2_bias, workspace.gemm2_output, workspace.gemm2_output_scale, + totalExpertsPerToken, gemm2HiddenSize, args.intermediate_size, totalLocalExperts, args.num_tokens, + workspace.num_non_exiting_ctas, workspace.total_num_padded_tokens, workspace.cta_idx_xy_to_batch_idx, + workspace.cta_idx_xy_to_mn_limit, workspace.bmm2_workspace, device, stream, config.gemm2Config, + args.valid_hidden_size.value_or(args.hidden_size), + args.valid_intermediate_size.value_or(args.intermediate_size)); + } // Run finalize if (args.do_finalize) diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.h b/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.h index 453a29671f4a..47e5a5b98e60 100644 --- a/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.h +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.h @@ -193,6 +193,17 @@ class Runner batchedGemm::trtllm::gen::Dtype dtypeRoutingLogits = batchedGemm::trtllm::gen::Dtype::Bfloat16, batchedGemm::trtllm::gen::Dtype dtypeRoutingBias = batchedGemm::trtllm::gen::Dtype::Bfloat16); + // run() for DeepSeekV3 routing without groups or fused shared experts, with the router GEMV in the same kernel: + // the routing logits are computed from routerHiddenStates [numTokens, hiddenDim] and routerWeight [numExperts, + // hiddenDim] (bf16) as dsv3_router_gemm_op computes them (see routingCustom::runFusedRouter, which checks that the + // shape is covered). + void runFusedRouter(void const* routerHiddenStates, void const* routerWeight, int32_t hiddenDim, + int32_t* routerWorkspace, void* routingBias, int32_t numTokens, int32_t numExperts, int32_t topK, + int32_t localExpertOffset, int32_t localNumExperts, float routedScalingFactor, int32_t* permutedIdxSize, + int32_t* expandedIdxToPermutedIdx, int32_t* permutedIdxToExpandedIdx, int32_t* permutedIdxToTokenIdx, + void* expertWeights, int32_t* ctaIdxXyToBatchIdx, int32_t* ctaIdxXyToMnLimit, int32_t* numNonExitingCtas, + cudaStream_t stream, batchedGemm::trtllm::gen::Dtype dtypeRoutingBias); + private: int32_t mTileTokensDim; }; @@ -288,6 +299,9 @@ struct MoERunnerArgs // [hidden_size/128, num_tokens] in float for e4m3 DS recipe // and [num_tokens, hidden_size/16] in float for e2m1 void* hidden_states_scale = nullptr; + // Optional: recorded on another stream once hidden_states and hidden_states_scale are written. The runner waits + // on it right before gemm1, their first reader, so the routing launched before gemm1 does not wait for them. + cudaEvent_t hidden_states_ready = nullptr; // Optional inputs: void* topk_weights = nullptr; // [num_tokens, top_k] with quantized weights @@ -352,6 +366,10 @@ struct MoERunnerArgs // finalize bool do_finalize{true}; + + // Run gemm2 on moe::dev::gemm2K128 instead of the Gemm2 cubin where it covers the shape (see Runner::run); the + // output rows are the cubin's, bitwise. + bool prefer_gemm2_k128{false}; }; struct MoEWorkspace @@ -441,6 +459,11 @@ class Runner // Kept so setOpsData can tell the activation launcher how much per-expert // tile padding the permuted row space carries. int32_t mTileTokensDim; + // Kept to tell whether gemm2K128 covers this runner's Gemm2 (see run()). + batchedGemm::trtllm::gen::Dtype mDtypeAct; + batchedGemm::trtllm::gen::Dtype mDtypeWeights; + bool mUseDeepSeekFp8; + bool mUseFineGrained; // This will be the cartesian product of the passing configs for gemm1 and gemm2 // This allows us to autotune the MoE as one operation instead of tuning gemm1 and gemm2 separately diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index 79cc1eb5dd7e..80fe9fc45c99 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -675,7 +675,10 @@ class TllmGenFmhaKernel if (mSM != kSM_107) { FmhaAutoTuner autoTuner(options, optionsFromArgs, multiProcessorCount); - return autoTuner.selectKernel(); + int32_t ctaDim = 512; + std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); + padTokenSparseMlaQTile(options); + return {options, optionsFromArgs, ctaDim}; } if (isGmemReduction(options.mMultiCtasKvMode) && options.mTileScheduler == TileScheduler::Static) @@ -692,10 +695,23 @@ class TllmGenFmhaKernel int32_t ctaDim = 512; FmhaAutoTuner autoTuner(options, optionsFromArgs, multiProcessorCount); std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); + padTokenSparseMlaQTile(options); limitCgaSmemReductionCtasKv(options); return {options, optionsFromArgs, ctaDim}; } + // The token-sparse MLA selection sets tileSizeQ to the Q heads per KV head, while the SwapsMmaAb kernels need a + // multiple of 8 Q-tile rows (SmemTile). With fewer than 8 heads per rank (for example 64 MLA heads over 16 ranks) + // pad the tile to 8, as the dense MLA selection does; the kernel masks the rows past the heads. + static void padTokenSparseMlaQTile(FmhaOptions& options) + { + if (options.mIsMlaGen && isTokenSparse(options.mSparseType) + && options.mFmhaKernelType == FmhaKernelType::SwapsMmaAbForGeneration && options.mTileSizeQ < 8) + { + options.mTileSizeQ = 8; + } + } + void limitCgaSmemReductionCtasKv(FmhaOptions& options) const { if (!isCgaSmemReduction(options.mMultiCtasKvMode)) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.cpp index d9bbee43305a..4a725a133284 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.cpp @@ -138,7 +138,7 @@ void KVCacheManagerV2UtilsBindings::initBindings(nb::module_& module) module.def( "copy_batch_block_offsets_to_device", [](at::Tensor input, at::Tensor output, at::Tensor copyIndex, at::Tensor indexScales, at::Tensor kvOffset, - uintptr_t stream) + uintptr_t stream, SizeType32 maxNumBlocks) { auto _input = from_torch(input); auto _output = from_torch(output); @@ -151,10 +151,11 @@ void KVCacheManagerV2UtilsBindings::initBindings(nb::module_& module) TLLM_CHECK_WITH_INFO(_indexScales.has_value(), "Invalid index scales tensor."); TLLM_CHECK_WITH_INFO(_kvOffset.has_value(), "Invalid kv offset tensor."); copyBatchBlockOffsetsToDevice(*(_input.value()), *(_output.value()), *(_copyIndex.value()), - *(_indexScales.value()), *(_kvOffset.value()), reinterpret_cast(stream)); + *(_indexScales.value()), *(_kvOffset.value()), maxNumBlocks, reinterpret_cast(stream)); }, nb::arg("input"), nb::arg("output"), nb::arg("copy_index"), nb::arg("index_scales"), nb::arg("kv_offset"), - nb::arg("stream"), nb::call_guard(), "Copy batch block indices to device"); + nb::arg("stream"), nb::arg("max_num_blocks") = -1, nb::call_guard(), + "Copy batch block indices to device; a non-negative max_num_blocks bounds the copied columns"); } } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index ab424318c862..c4dd5613839e 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -74,7 +74,9 @@ add_th_common_family( dsv3RopeOp.cpp convertReqIndexToGlobalOp.cpp nvfp4MlaKvCacheGatherOp.cpp - trtllmGenQKVProcessOp.cpp) + trtllmGenQKVProcessOp.cpp + glmMlaKvHalfAppendOp.cpp + glmSparseMlaFlashDecodeOp.cpp) add_th_common_family( th_common_gemm_quant @@ -98,6 +100,8 @@ add_th_common_family( fp8RowwiseGemm.cpp fp8Quantize.cpp dsv3FusedAGemmOp.cpp + glmDenseMlpSmallMOp.cpp + glmFusedAGemmOp.cpp deepseekV4QNormOp.cpp fusedAddRMSNormQuant.cpp fusedActivationQuant.cpp @@ -108,6 +112,7 @@ add_th_common_family( mxfp8Gemm.cpp fusedCatFp4Op.cpp fusedCatFp8Op.cpp + glmFusedEmbeddingOp.cpp weightOnlyQuantGemm.cpp weightOnlyQuantOp.cpp finegrained_mixed_dtype_gemm_thop.cpp @@ -116,8 +121,14 @@ add_th_common_family( # TODO: th_common_moe does not compile with USING_OSS_CUTLASS_MOE_GEMM=OFF # because moe/utils/moeUtilOp.cpp depends on OSS-specific MoE utility types and # function signatures. -add_th_common_family(th_common_moe fusedTopkSoftmax.cpp llama4MinLatency.cpp - noAuxTcOp.cpp dsv3RouterGemmOp.cpp) +add_th_common_family( + th_common_moe + fusedTopkSoftmax.cpp + llama4MinLatency.cpp + noAuxTcOp.cpp + dsv3RouterGemmOp.cpp + glmRouterGemmOp.cpp + glmLowMBf16MoeOp.cpp) # MoE torch ops live in moe//; that directory adds them to # th_common_moe with target_sources() so the list stays next to the files. @@ -128,7 +139,10 @@ add_th_common_family( allgatherOp.cpp allreduceOp.cpp alltoallOp.cpp + glmDsaHistTopKOp.cpp asyncUlyssesOp.cpp + glmMnnvlAllgatherOp.cpp + glmMnnvlHopOp.cpp ulyssesPostUnscatterOp.cpp ulyssesPermuteScatterOp.cpp helixPostProcessOp.cpp @@ -139,7 +153,8 @@ add_th_common_family( userbuffersFinalizeOp.cpp userbuffersTensor.cpp virtualMemoryAllocator.cpp - fusedGemmAllreduceOp.cpp) + fusedGemmAllreduceOp.cpp + glmDraftSampleGlueOp.cpp) add_th_common_family(th_common_sequence causalConv1dOp.cpp mambaConv1dOp.cpp mamba2MTPSSMCacheOp.cpp selectiveScanOp.cpp) @@ -159,7 +174,8 @@ add_th_common_family( loraOp.cpp inplaceSliceCopyOp.cpp mhcOp.cpp - compressorOp.cpp) + compressorOp.cpp + glmVerifyHeadAcceptOp.cpp) add_library( th_common SHARED diff --git a/cpp/tensorrt_llm/thop/glmDenseMlpSmallMOp.cpp b/cpp/tensorrt_llm/thop/glmDenseMlpSmallMOp.cpp new file mode 100644 index 000000000000..91b36a423ec4 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmDenseMlpSmallMOp.cpp @@ -0,0 +1,115 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace tkd = tensorrt_llm::kernels::glm_moe_dsa::dense_mlp_small_m; + +//! The kernels use clusters, bulk copies and st.async; below sm_90 their bodies compile to nothing. +constexpr int kMinSmVersion = 90; + +void checkBf16Matrix(at::Tensor const& t, char const* name) +{ + CHECK_INPUT(t, at::ScalarType::BFloat16); + TORCH_CHECK(t.dim() == 2, name, " must be 2-D"); + TORCH_CHECK(reinterpret_cast(t.data_ptr()) % 16 == 0, name, " must be 16-byte aligned"); +} + +} // namespace + +//! Whether the kernels serve M rows of a dense MLP with this hidden and rank-local intermediate size (host only). +bool glmDenseMlpSmallMSupported(int64_t numRows, int64_t hidden, int64_t inter) +{ + return tkd::isGateUpSupported(static_cast(numRows), static_cast(hidden), static_cast(inter)) + && tkd::isDownSupported(static_cast(numRows), static_cast(inter), static_cast(hidden)); +} + +//! The rank-local partial of a BF16 gated MLP with plain SwiGLU (no output all-reduce): x [M, H], +//! gate_up_weight [2 I, H] (gate rows first), down_weight [H, I] -> [M, H]. gate_up_tiling / down_tiling select a +//! compiled tiling (-1: the default of the shape). +at::Tensor glmDenseMlpSmallM(at::Tensor const& x, at::Tensor const& gateUpWeight, at::Tensor const& downWeight, + int64_t gateUpTiling, int64_t downTiling) +{ + checkBf16Matrix(x, "x"); + checkBf16Matrix(gateUpWeight, "gate_up_weight"); + checkBf16Matrix(downWeight, "down_weight"); + TORCH_CHECK(gateUpWeight.device() == x.device() && downWeight.device() == x.device(), + "glm_dense_mlp_small_m: all tensors must be on the device of x"); + TORCH_CHECK(tensorrt_llm::common::getSMVersion() >= kMinSmVersion, "glm_dense_mlp_small_m needs sm_90 or newer"); + int64_t const numRows = x.size(0); + int64_t const hidden = x.size(1); + int64_t const inter = downWeight.size(1); + TORCH_CHECK(gateUpWeight.size(0) == 2 * inter && gateUpWeight.size(1) == hidden && downWeight.size(0) == hidden, + "glm_dense_mlp_small_m: gate_up_weight ", gateUpWeight.sizes(), " and down_weight ", downWeight.sizes(), + " do not match x ", x.sizes()); + TORCH_CHECK(tkd::isGateUpSupported(static_cast(numRows), static_cast(hidden), static_cast(inter), + static_cast(gateUpTiling)) + && tkd::isDownSupported(static_cast(numRows), static_cast(inter), static_cast(hidden), + static_cast(downTiling)), + "glm_dense_mlp_small_m: unsupported shape (M ", numRows, ", hidden ", hidden, ", inter ", inter, ") or tiling"); + + at::Tensor h = at::empty({numRows, inter}, x.options()); + at::Tensor out = at::empty({numRows, hidden}, x.options()); + auto const stream = at::cuda::getCurrentCUDAStream(x.get_device()); + bool const enablePdl = tensorrt_llm::common::getEnvEnablePDL(); + TLLM_CUDA_CHECK(tkd::launchGateUpSwiglu(static_cast<__nv_bfloat16 const*>(x.data_ptr()), + static_cast<__nv_bfloat16 const*>(gateUpWeight.data_ptr()), static_cast<__nv_bfloat16*>(h.data_ptr()), + static_cast(numRows), static_cast(hidden), static_cast(inter), static_cast(gateUpTiling), + enablePdl, stream)); + TLLM_CUDA_CHECK(tkd::launchDown(static_cast<__nv_bfloat16 const*>(h.data_ptr()), + static_cast<__nv_bfloat16 const*>(downWeight.data_ptr()), static_cast<__nv_bfloat16*>(out.data_ptr()), + static_cast(numRows), static_cast(inter), static_cast(hidden), static_cast(downTiling), + enablePdl, stream)); + return out; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("glm_dense_mlp_small_m_supported(int num_rows, int hidden, int inter) -> bool"); + m.def( + "glm_dense_mlp_small_m(Tensor x, Tensor gate_up_weight, Tensor down_weight, int gate_up_tiling=-1, " + "int down_tiling=-1) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CompositeExplicitAutograd, m) +{ + m.impl("glm_dense_mlp_small_m_supported", &tensorrt_llm::torch_ext::glmDenseMlpSmallMSupported); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_dense_mlp_small_m", &tensorrt_llm::torch_ext::glmDenseMlpSmallM); +} diff --git a/cpp/tensorrt_llm/thop/glmDraftSampleGlueOp.cpp b/cpp/tensorrt_llm/thop/glmDraftSampleGlueOp.cpp new file mode 100644 index 000000000000..b9bbcadff112 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmDraftSampleGlueOp.cpp @@ -0,0 +1,366 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/mcastDevMemUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.h" +#include "tensorrt_llm/runtime/mcastDeviceMemory.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace tk = tensorrt_llm::kernels::glm_moe_dsa; + +bool isAligned16(at::Tensor const& t) +{ + return reinterpret_cast(t.data_ptr()) % 16 == 0; +} + +// The logits, exchange workspace and draft-token fields of both draft-step ops; slotBytes is what one rank publishes +// per row. Returns the number of rows. +int64_t setDraftStepExchange(tk::DraftSampleGlueParams& params, char const* op, at::Tensor const& logits, + at::Tensor& exchange_buffer, at::Tensor& exchange_flags, int64_t group_size, at::Tensor& draft_tokens, + int64_t draft_step, int64_t slotBytes) +{ + auto* mcastMem = tensorrt_llm::common::findMcastDevMemBuffer(exchange_buffer.data_ptr()); + TORCH_CHECK(mcastMem != nullptr, op, ": exchange_buffer must come from a McastGPUBuffer"); + TORCH_CHECK(static_cast(mcastMem->getWorldSize()) == group_size, op, ": group_size ", group_size, + " does not match the workspace group of ", mcastMem->getWorldSize(), " ranks"); + TORCH_CHECK(exchange_flags.element_size() == 4 && exchange_flags.numel() >= 9, op, + ": exchange_flags must hold the nine 32-bit Lamport flags"); + CHECK_INPUT(logits, at::ScalarType::BFloat16); + TORCH_CHECK(logits.dim() == 2 && logits.size(0) >= 1 && logits.size(0) <= tk::kDraftSampleGlueMaxRows, op, + ": logits must be [rows, vocab_shard] with 1 to ", tk::kDraftSampleGlueMaxRows, " rows"); + int64_t const rows = logits.size(0); + TORCH_CHECK(logits.size(1) % 8 == 0 && isAligned16(logits), op, + ": logits rows must be 16-byte aligned and a multiple of 8 columns"); + int64_t const bufferBytes = exchange_buffer.numel() * exchange_buffer.element_size() / 3; + TORCH_CHECK( + bufferBytes >= rows * group_size * slotBytes, op, ": the Lamport buffer is too small for ", rows, " rows"); + CHECK_INPUT(draft_tokens, at::ScalarType::Int); + TORCH_CHECK( + draft_tokens.dim() == 2 && draft_tokens.size(0) == rows && draft_step >= 0 && draft_step < draft_tokens.size(1), + op, ": draft_tokens must be [rows, draft_len] with draft_step inside"); + + params.logits = reinterpret_cast<__nv_bfloat16 const*>(logits.data_ptr()); + params.logitsStride = logits.stride(0); + params.numRows = static_cast(rows); + params.vocabShard = static_cast(logits.size(1)); + params.rank = static_cast(mcastMem->getRank()); + params.tpSize = static_cast(group_size); + params.exchangeLocal = exchange_buffer.data_ptr(); + params.exchangeMcast = mcastMem->getMulticastPtr(); + params.exchangeFlags = reinterpret_cast(exchange_flags.data_ptr()); + params.draftTokens = draft_tokens.data_ptr(); + params.numDraftTokens = static_cast(draft_tokens.size(1)); + params.draftStep = static_cast(draft_step); + return rows; +} + +// The next step's position ids, sequence / KV lengths and DSA decode metadata of both draft-step ops. +void setDraftStepState(tk::DraftSampleGlueParams& params, char const* op, int64_t rows, + std::optional const& positions_in, std::optional& positions_out, + std::optional const& gather_ids, std::optional& seq_lens, bool fill_seq_lens, + std::optional& kv_lens, int64_t kv_update, std::optional const& rewind_accepted, + int64_t draft_len, std::optional const& block_offsets, std::vector const& dsa_layout, + std::optional& req_idx_per_token, std::optional& slot_mapping_fp8, + std::optional& slot_mapping_scale, std::optional& gen_kv_indptr, + std::optional& gen_cached_token_indptr) +{ + if (positions_out.has_value()) + { + TORCH_CHECK(positions_in.has_value(), op, ": positions_out needs positions_in"); + auto const posType = positions_in->scalar_type(); + TORCH_CHECK(posType == at::ScalarType::Int || posType == at::ScalarType::Long, op, + ": position ids must be int32 or int64"); + CHECK_INPUT(positions_in.value(), posType); + CHECK_INPUT(positions_out.value(), posType); + TORCH_CHECK(positions_out->numel() >= rows, op, ": positions_out holds one id per row"); + params.positionsIn = positions_in->data_ptr(); + params.positionsOut = positions_out->data_ptr(); + params.positions64 = posType == at::ScalarType::Long; + if (gather_ids.has_value()) + { + CHECK_INPUT(gather_ids.value(), at::ScalarType::Long); + TORCH_CHECK(gather_ids->numel() >= rows, op, ": gather_ids holds one index per row"); + params.gatherIds = gather_ids->data_ptr(); + } + } + + if (seq_lens.has_value()) + { + CHECK_INPUT(seq_lens.value(), at::ScalarType::Int); + TORCH_CHECK(seq_lens->numel() >= rows, op, ": seq_lens holds one length per row"); + params.seqLens = seq_lens->data_ptr(); + params.fillSeqLens = fill_seq_lens; + } + if (kv_lens.has_value()) + { + CHECK_INPUT(kv_lens.value(), at::ScalarType::Int); + TORCH_CHECK(kv_lens->numel() >= rows, op, ": kv_lens holds one length per row"); + TORCH_CHECK(kv_update >= 0 && kv_update <= 2, op, ": unknown kv_update ", kv_update); + params.kvLens = kv_lens->data_ptr(); + params.kvUpdate = static_cast(kv_update); + if (params.kvUpdate == tk::DraftSampleGlueKvUpdate::kRewind) + { + TORCH_CHECK(rewind_accepted.has_value(), op, ": kv_update 2 needs rewind_accepted"); + CHECK_INPUT(rewind_accepted.value(), at::ScalarType::Int); + TORCH_CHECK(rewind_accepted->numel() >= rows, op, ": rewind_accepted holds one count per row"); + params.kvRewindAccepted = rewind_accepted->data_ptr(); + } + params.draftLen = static_cast(draft_len); + } + + if (req_idx_per_token.has_value()) + { + TORCH_CHECK(seq_lens.has_value() && kv_lens.has_value() && block_offsets.has_value() + && slot_mapping_fp8.has_value() && slot_mapping_scale.has_value() && gen_kv_indptr.has_value() + && gen_cached_token_indptr.has_value() && dsa_layout.size() == 7, + op, ": the DSA metadata needs the lengths, the block offsets, all five outputs and the layout"); + CHECK_TH_CUDA(block_offsets.value()); + TORCH_CHECK(block_offsets->scalar_type() == at::ScalarType::Int && block_offsets->dim() == 2 + && block_offsets->size(0) >= rows, + op, ": block_offsets must be int32 [rows, max_blocks]"); + int64_t const numTokens = dsa_layout[0]; + CHECK_INPUT(req_idx_per_token.value(), at::ScalarType::Int); + CHECK_INPUT(slot_mapping_fp8.value(), at::ScalarType::Long); + CHECK_INPUT(slot_mapping_scale.value(), at::ScalarType::Long); + CHECK_INPUT(gen_kv_indptr.value(), at::ScalarType::Long); + CHECK_INPUT(gen_cached_token_indptr.value(), at::ScalarType::Long); + TORCH_CHECK(req_idx_per_token->numel() >= numTokens && slot_mapping_fp8->numel() >= numTokens + && slot_mapping_scale->numel() >= numTokens && gen_kv_indptr->numel() >= rows + 1 + && gen_cached_token_indptr->numel() >= rows + 1, + op, ": DSA metadata outputs too small"); + params.blockOffsets = block_offsets->data_ptr(); + params.blockOffsetsStride0 = block_offsets->stride(0); + params.blockOffsetsStride1 = block_offsets->stride(1); + params.maxBlocks = static_cast(block_offsets->size(1)); + params.numTokens = static_cast(numTokens); + params.maxQueryLanes = static_cast(dsa_layout[1]); + params.tokensPerBlock = static_cast(dsa_layout[2]); + params.dataBytesPerToken = dsa_layout[3]; + params.scaleSize = dsa_layout[4]; + params.blockStride = dsa_layout[5]; + params.scaleBaseOffset = dsa_layout[6]; + params.reqIdxPerToken = req_idx_per_token->data_ptr(); + params.slotMappingFp8 = slot_mapping_fp8->data_ptr(); + params.slotMappingScale = slot_mapping_scale->data_ptr(); + params.genKvIndptr = gen_kv_indptr->data_ptr(); + params.genCachedTokenIndptr = gen_cached_token_indptr->data_ptr(); + } +} + +} // namespace + +// One MTP draft step after the vocabulary-sharded draft LM head (see draftSampleGlue.h): the greedy draft token of +// every row into column draft_step of draft_tokens, through the dedicated MNNVL Lamport workspace exchange_buffer / +// exchange_flags; optionally the next step's eh_proj input window (glue_out), position ids, sequence and KV lengths +// (kv_update 0: none, 1: + 1, 2: - (draft_len - rewind_accepted)), the DSA decode metadata (dsa_layout = [num_tokens, +// max_query_lanes, tokens_per_block, data_bytes_per_token, scale_size, block_stride, scale_base_offset]) and the +// next-new-tokens rows after the last step. +void glmDraftSampleGlue(at::Tensor const& logits, at::Tensor& exchange_buffer, at::Tensor& exchange_flags, + int64_t group_size, at::Tensor& draft_tokens, int64_t draft_step, std::optional const& embed_table, + std::optional const& hidden, std::optional const& enorm_weight, + std::optional const& hnorm_weight, double eps, int64_t col_lo, int64_t col_hi, + std::optional glue_out, std::optional const& positions_in, + std::optional positions_out, std::optional const& gather_ids, + std::optional seq_lens, bool fill_seq_lens, std::optional kv_lens, int64_t kv_update, + std::optional const& rewind_accepted, int64_t draft_len, std::optional const& block_offsets, + std::vector const& dsa_layout, std::optional req_idx_per_token, + std::optional slot_mapping_fp8, std::optional slot_mapping_scale, + std::optional gen_kv_indptr, std::optional gen_cached_token_indptr, + std::optional const& accepted_tokens, std::optional const& batch_indices, + std::optional const& num_accepted, std::optional next_new_tokens) +{ + tk::DraftSampleGlueParams params; + int64_t const rows = setDraftStepExchange(params, "glm_draft_sample_glue", logits, exchange_buffer, exchange_flags, + group_size, draft_tokens, draft_step, tk::kDraftSampleGluePacketBytes); + + if (glue_out.has_value()) + { + constexpr int64_t kHidden = tk::kDraftSampleGlueHidden; + TORCH_CHECK(col_lo >= 0 && col_lo < col_hi && col_hi <= 2 * kHidden && col_lo % 8 == 0 && col_hi % 8 == 0, + "glm_draft_sample_glue: invalid eh_proj input window [", col_lo, ", ", col_hi, ")"); + TORCH_CHECK( + enorm_weight.has_value() && hnorm_weight.has_value(), "glm_draft_sample_glue: norm weights missing"); + for (at::Tensor const& weight : {enorm_weight.value(), hnorm_weight.value()}) + { + CHECK_INPUT(weight, at::ScalarType::BFloat16); + TORCH_CHECK( + weight.numel() == kHidden && isAligned16(weight), "glm_draft_sample_glue: norm weights must be [6144]"); + } + TORCH_CHECK(glue_out->scalar_type() == at::ScalarType::BFloat16 && glue_out->dim() == 2 + && glue_out->size(0) == rows && glue_out->size(1) == col_hi - col_lo && glue_out->stride(1) == 1 + && glue_out->stride(0) % 8 == 0 && isAligned16(*glue_out), + "glm_draft_sample_glue: glue_out must be a 16-byte aligned BF16 [rows, col_hi - col_lo] view"); + if (col_lo < kHidden) + { + TORCH_CHECK(embed_table.has_value(), "glm_draft_sample_glue: the window needs the embedding table"); + CHECK_INPUT(embed_table.value(), at::ScalarType::BFloat16); + TORCH_CHECK(embed_table->dim() == 2 && embed_table->size(1) == kHidden && isAligned16(*embed_table), + "glm_draft_sample_glue: embed_table must be [vocab, 6144]"); + params.embedTable = reinterpret_cast<__nv_bfloat16 const*>(embed_table->data_ptr()); + params.vocabSize = embed_table->size(0); + } + if (col_hi > kHidden) + { + TORCH_CHECK(hidden.has_value(), "glm_draft_sample_glue: the window needs the hidden rows"); + TORCH_CHECK(hidden->scalar_type() == at::ScalarType::BFloat16 && hidden->dim() == 2 + && hidden->size(0) == rows && hidden->size(1) == kHidden && hidden->stride(1) == 1 + && hidden->stride(0) % 8 == 0 && isAligned16(*hidden), + "glm_draft_sample_glue: hidden must be a 16-byte aligned BF16 [rows, 6144] view"); + params.hidden = reinterpret_cast<__nv_bfloat16 const*>(hidden->data_ptr()); + params.hiddenStride = hidden->stride(0); + } + params.enormWeight = reinterpret_cast<__nv_bfloat16 const*>(enorm_weight->data_ptr()); + params.hnormWeight = reinterpret_cast<__nv_bfloat16 const*>(hnorm_weight->data_ptr()); + params.eps = static_cast(eps); + params.colLo = static_cast(col_lo); + params.colHi = static_cast(col_hi); + params.glueOut = reinterpret_cast<__nv_bfloat16*>(glue_out->data_ptr()); + params.glueOutStride = glue_out->stride(0); + } + + setDraftStepState(params, "glm_draft_sample_glue", rows, positions_in, positions_out, gather_ids, seq_lens, + fill_seq_lens, kv_lens, kv_update, rewind_accepted, draft_len, block_offsets, dsa_layout, req_idx_per_token, + slot_mapping_fp8, slot_mapping_scale, gen_kv_indptr, gen_cached_token_indptr); + + if (next_new_tokens.has_value()) + { + TORCH_CHECK(accepted_tokens.has_value() && batch_indices.has_value() && num_accepted.has_value(), + "glm_draft_sample_glue: next_new_tokens needs the accepted tokens, batch indices and counts"); + CHECK_INPUT(next_new_tokens.value(), at::ScalarType::Int); + CHECK_TH_CUDA(accepted_tokens.value()); + TORCH_CHECK(accepted_tokens->scalar_type() == at::ScalarType::Int && accepted_tokens->dim() == 2 + && accepted_tokens->stride(1) == 1, + "glm_draft_sample_glue: accepted_tokens must be int32 [batch, draft_len + 1] rows"); + CHECK_INPUT(batch_indices.value(), at::ScalarType::Int); + CHECK_INPUT(num_accepted.value(), at::ScalarType::Int); + TORCH_CHECK(next_new_tokens->numel() >= rows * (draft_tokens.size(1) + 1) && batch_indices->numel() >= rows + && num_accepted->numel() >= rows, + "glm_draft_sample_glue: next_new_tokens must be [rows, draft_len + 1]"); + params.acceptedTokens = accepted_tokens->data_ptr(); + params.acceptedTokensStride = accepted_tokens->stride(0); + params.batchIndices = batch_indices->data_ptr(); + params.numAccepted = num_accepted->data_ptr(); + params.nextNewTokens = next_new_tokens->data_ptr(); + } + + tk::invokeDraftSampleGlue(params, at::cuda::getCurrentCUDAStream(logits.get_device())); +} + +// One MTP draft step that also replaces the next step's eh_proj with its all-reduce and input_layernorm (see +// invokeDraftEhprojTable): the draft token and the next step's state as glm_draft_sample_glue computes them, +// residual_out = x = BF16(token_table[token] + hidden_half), one FP32 add, and normed_out = input_layernorm(x) with +// norm_weight and eps. hidden_half is this rank's FP32 [rows, 6144 / group_size] share of eh_proj's hidden half, +// token_table this rank's [vocab_shard, 6144] rows; exchange_buffer / exchange_flags are a dedicated MNNVL Lamport +// workspace of draftEhprojTableSlotBytes(group_size) bytes per rank and row. +void glmDraftEhprojTable(at::Tensor const& logits, at::Tensor& exchange_buffer, at::Tensor& exchange_flags, + int64_t group_size, at::Tensor& draft_tokens, int64_t draft_step, at::Tensor const& hidden_half, + at::Tensor const& token_table, at::Tensor const& norm_weight, double eps, at::Tensor& normed_out, + at::Tensor& residual_out, std::optional const& positions_in, std::optional positions_out, + std::optional const& gather_ids, std::optional seq_lens, bool fill_seq_lens, + std::optional kv_lens, int64_t kv_update, std::optional const& rewind_accepted, + int64_t draft_len, std::optional const& block_offsets, std::vector const& dsa_layout, + std::optional req_idx_per_token, std::optional slot_mapping_fp8, + std::optional slot_mapping_scale, std::optional gen_kv_indptr, + std::optional gen_cached_token_indptr) +{ + constexpr int64_t kHidden = tk::kDraftSampleGlueHidden; + TORCH_CHECK(tk::draftEhprojTableSupportsTpSize(static_cast(group_size)), + "glm_draft_ehproj_table: group_size must be 2 to ", tk::kDraftEhprojTableMaxRanks, + " ranks with whole 8-column shares of the hidden size, got ", group_size); + tk::DraftEhprojTableParams params; + int64_t const rows = setDraftStepExchange(params.step, "glm_draft_ehproj_table", logits, exchange_buffer, + exchange_flags, group_size, draft_tokens, draft_step, tk::draftEhprojTableSlotBytes(group_size)); + + TORCH_CHECK(hidden_half.is_cuda() && hidden_half.scalar_type() == at::ScalarType::Float && hidden_half.dim() == 2 + && hidden_half.size(0) == rows && hidden_half.size(1) == kHidden / group_size && hidden_half.stride(1) == 1 + && hidden_half.stride(0) % 8 == 0 && reinterpret_cast(hidden_half.data_ptr()) % 32 == 0, + "glm_draft_ehproj_table: hidden_half must be a 32-byte aligned FP32 [rows, 6144 / group_size] view"); + CHECK_INPUT(token_table, at::ScalarType::BFloat16); + TORCH_CHECK(token_table.dim() == 2 && token_table.size(0) == logits.size(1) && token_table.size(1) == kHidden + && isAligned16(token_table), + "glm_draft_ehproj_table: token_table must be [vocab_shard, 6144]"); + CHECK_INPUT(norm_weight, at::ScalarType::BFloat16); + TORCH_CHECK(norm_weight.numel() == kHidden && isAligned16(norm_weight), + "glm_draft_ehproj_table: norm_weight must be [6144]"); + for (at::Tensor const& out : {normed_out, residual_out}) + { + TORCH_CHECK(out.is_cuda() && out.scalar_type() == at::ScalarType::BFloat16 && out.dim() == 2 + && out.size(0) == rows && out.size(1) == kHidden && out.stride(1) == 1 + && out.stride(0) == normed_out.stride(0) && out.stride(0) % 8 == 0 && isAligned16(out), + "glm_draft_ehproj_table: normed_out and residual_out must be 16-byte aligned BF16 [rows, 6144] views with " + "one row stride"); + } + params.hiddenHalf = hidden_half.data_ptr(); + params.hiddenHalfStride = hidden_half.stride(0); + params.tokenTable = reinterpret_cast<__nv_bfloat16 const*>(token_table.data_ptr()); + params.normWeight = reinterpret_cast<__nv_bfloat16 const*>(norm_weight.data_ptr()); + params.normEps = static_cast(eps); + params.normedOut = reinterpret_cast<__nv_bfloat16*>(normed_out.data_ptr()); + params.residualOut = reinterpret_cast<__nv_bfloat16*>(residual_out.data_ptr()); + params.outStride = normed_out.stride(0); + + setDraftStepState(params.step, "glm_draft_ehproj_table", rows, positions_in, positions_out, gather_ids, seq_lens, + fill_seq_lens, kv_lens, kv_update, rewind_accepted, draft_len, block_offsets, dsa_layout, req_idx_per_token, + slot_mapping_fp8, slot_mapping_scale, gen_kv_indptr, gen_cached_token_indptr); + + tk::invokeDraftEhprojTable(params, at::cuda::getCurrentCUDAStream(logits.get_device())); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_draft_sample_glue(Tensor logits, Tensor(a!) exchange_buffer, Tensor(b!) exchange_flags, int group_size, " + "Tensor(c!) draft_tokens, int draft_step, Tensor? embed_table, Tensor? hidden, Tensor? enorm_weight, " + "Tensor? hnorm_weight, float eps, int col_lo, int col_hi, Tensor(d!)? glue_out, Tensor? positions_in, " + "Tensor(e!)? positions_out, Tensor? gather_ids, Tensor(f!)? seq_lens, bool fill_seq_lens, " + "Tensor(g!)? kv_lens, int kv_update, Tensor? rewind_accepted, int draft_len, Tensor? block_offsets, " + "int[] dsa_layout, Tensor(h!)? req_idx_per_token, Tensor(i!)? slot_mapping_fp8, " + "Tensor(j!)? slot_mapping_scale, Tensor(k!)? gen_kv_indptr, Tensor(l!)? gen_cached_token_indptr, " + "Tensor? accepted_tokens, Tensor? batch_indices, Tensor? num_accepted, Tensor(m!)? next_new_tokens) -> ()"); + m.def( + "glm_draft_ehproj_table(Tensor logits, Tensor(a!) exchange_buffer, Tensor(b!) exchange_flags, int group_size, " + "Tensor(c!) draft_tokens, int draft_step, Tensor hidden_half, Tensor token_table, Tensor norm_weight, " + "float eps, Tensor(d!) normed_out, Tensor(e!) residual_out, Tensor? positions_in, Tensor(f!)? positions_out, " + "Tensor? gather_ids, Tensor(g!)? seq_lens, bool fill_seq_lens, Tensor(h!)? kv_lens, int kv_update, " + "Tensor? rewind_accepted, int draft_len, Tensor? block_offsets, int[] dsa_layout, " + "Tensor(i!)? req_idx_per_token, Tensor(j!)? slot_mapping_fp8, Tensor(k!)? slot_mapping_scale, " + "Tensor(l!)? gen_kv_indptr, Tensor(m!)? gen_cached_token_indptr) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_draft_sample_glue", &tensorrt_llm::torch_ext::glmDraftSampleGlue); + m.impl("glm_draft_ehproj_table", &tensorrt_llm::torch_ext::glmDraftEhprojTable); +} diff --git a/cpp/tensorrt_llm/thop/glmDsaHistTopKOp.cpp b/cpp/tensorrt_llm/thop/glmDsaHistTopKOp.cpp new file mode 100644 index 000000000000..fb0ab897d100 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmDsaHistTopKOp.cpp @@ -0,0 +1,120 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/mcastDevMemUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.h" +#include "tensorrt_llm/runtime/mcastDeviceMemory.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include + +#include + +namespace tkh = tensorrt_llm::kernels::glm_moe_dsa::hist_topk; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +// Exact top-2048 of one generation request's key-sharded GLM DSA indexer logits (glmDsaHistTopK.h): this rank's shard +// logits [num_rows, >= shard size] FP32 in, the request-local token indices [num_rows, 2048] int32 out, identical on +// every rank of the group. comm_buffer is this rank's unicast view of a McastGPUBuffer workspace of +// kWorkspaceWords int32 words, all -1 at creation; state holds kStateWords int32, zero at creation; both are private +// to this op. +void glmDsaHistTopK(at::Tensor const& logits, at::Tensor const& kvLens, at::Tensor& out, at::Tensor& commBuffer, + at::Tensor& state, int64_t nextN, int64_t groupSize) +{ + auto* mcastMem = tensorrt_llm::common::findMcastDevMemBuffer(commBuffer.data_ptr()); + TORCH_CHECK(mcastMem != nullptr, "glm_dsa_hist_topk: comm_buffer must come from a McastGPUBuffer"); + TORCH_CHECK(static_cast(mcastMem->getWorldSize()) == groupSize, "glm_dsa_hist_topk: group_size ", + groupSize, " does not match the workspace group of ", mcastMem->getWorldSize(), " ranks"); + CHECK_TH_CUDA(logits); + CHECK_TH_CUDA(kvLens); + CHECK_TH_CUDA(out); + CHECK_TH_CUDA(state); + TORCH_CHECK(logits.scalar_type() == at::ScalarType::Float && logits.dim() == 2 && logits.stride(1) == 1, + "glm_dsa_hist_topk: logits must be a row-major FP32 matrix"); + TORCH_CHECK(kvLens.scalar_type() == at::ScalarType::Int && kvLens.numel() >= 1, "glm_dsa_hist_topk: kv_lens int32"); + TORCH_CHECK(out.scalar_type() == at::ScalarType::Int && out.dim() == 2 && out.size(1) == tkh::kTopK + && out.stride(1) == 1 && out.stride(0) == tkh::kTopK, + "glm_dsa_hist_topk: out must be a contiguous int32 [num_rows, 2048] matrix"); + int64_t const numRows = out.size(0); + TORCH_CHECK(numRows >= 1 && numRows <= tkh::kMaxRows && logits.size(0) >= numRows && nextN >= numRows, + "glm_dsa_hist_topk: 1..8 rows of one request"); + TORCH_CHECK(logits.size(1) <= tkh::kMaxShard, "glm_dsa_hist_topk: at most ", tkh::kMaxShard, " shard columns"); + TORCH_CHECK(commBuffer.scalar_type() == at::ScalarType::Int && commBuffer.numel() >= tkh::kWorkspaceWords, + "glm_dsa_hist_topk: comm_buffer must hold ", tkh::kWorkspaceWords, " int32 words"); + TORCH_CHECK(state.scalar_type() == at::ScalarType::Int && state.numel() >= tkh::kStateWords, + "glm_dsa_hist_topk: state must hold ", tkh::kStateWords, " int32 words"); + + tkh::HistTopKParams params; + params.logits = logits.data_ptr(); + params.logitsStride = logits.stride(0); + params.localCapacity = static_cast(logits.size(1)); + params.kvLens = kvLens.data_ptr(); + params.out = out.data_ptr(); + params.numRows = static_cast(numRows); + params.nextN = static_cast(nextN); + params.rank = static_cast(mcastMem->getRank()); + params.worldSize = static_cast(groupSize); + params.mcWorkspace = static_cast(mcastMem->getMulticastPtr()); + params.ucWorkspace = reinterpret_cast(commBuffer.data_ptr()); + params.state = state.data_ptr(); + tkh::invokeDsaHistTopK(params, at::cuda::getCurrentCUDAStream(logits.get_device())); +} + +// Shard metadata of the generation request for this rank (glmDsaHistTopK.h): shard_ctx [1] and shard_block_table +// [1, width] int32, from the request's KV length and block table row. +void glmDsaHistShardMeta(at::Tensor const& kvLens, at::Tensor const& blockTable, at::Tensor& shardCtx, + at::Tensor& shardBlockTable, int64_t groupSize, int64_t rank, int64_t tokensPerBlock) +{ + CHECK_TH_CUDA(kvLens); + CHECK_TH_CUDA(blockTable); + TORCH_CHECK(kvLens.scalar_type() == at::ScalarType::Int && kvLens.numel() >= 1, + "glm_dsa_hist_shard_meta: kv_lens must be int32"); + TORCH_CHECK(blockTable.scalar_type() == at::ScalarType::Int && blockTable.dim() >= 1 && blockTable.stride(-1) == 1, + "glm_dsa_hist_shard_meta: block_table must be int32 with contiguous rows"); + TORCH_CHECK(shardCtx.scalar_type() == at::ScalarType::Int && shardCtx.numel() >= 1, + "glm_dsa_hist_shard_meta: shard_ctx must be int32"); + TORCH_CHECK(shardBlockTable.scalar_type() == at::ScalarType::Int && shardBlockTable.is_contiguous(), + "glm_dsa_hist_shard_meta: shard_block_table must be contiguous int32"); + tkh::invokeDsaHistShardMeta(kvLens.data_ptr(), blockTable.data_ptr(), + static_cast(blockTable.size(-1)), static_cast(groupSize), static_cast(rank), + static_cast(tokensPerBlock), shardCtx.data_ptr(), shardBlockTable.data_ptr(), + static_cast(shardBlockTable.numel()), at::cuda::getCurrentCUDAStream(kvLens.get_device())); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_dsa_hist_topk(Tensor logits, Tensor kv_lens, Tensor(a!) out, Tensor(b!) comm_buffer, Tensor(c!) state, " + "int next_n, int group_size) -> ()"); + m.def( + "glm_dsa_hist_shard_meta(Tensor kv_lens, Tensor block_table, Tensor(a!) shard_ctx, " + "Tensor(b!) shard_block_table, int group_size, int rank, int tokens_per_block) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_dsa_hist_topk", &tensorrt_llm::torch_ext::glmDsaHistTopK); + m.impl("glm_dsa_hist_shard_meta", &tensorrt_llm::torch_ext::glmDsaHistShardMeta); +} diff --git a/cpp/tensorrt_llm/thop/glmFusedAGemmOp.cpp b/cpp/tensorrt_llm/thop/glmFusedAGemmOp.cpp new file mode 100644 index 000000000000..2a33b5fc1e1f --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmFusedAGemmOp.cpp @@ -0,0 +1,130 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.h" + +#include +#include +#include +#include + +namespace tk = tensorrt_llm::kernels; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +using tk::glm_moe_dsa::FusedAL2Prefetch; + +constexpr int64_t kTmaAlignmentBytes = 16; + +bool isAligned(torch::Tensor const& tensor) +{ + return reinterpret_cast(tensor.data_ptr()) % kTmaAlignmentBytes == 0; +} + +// One L2 prefetch range per tensor: its contiguous bytes on the GEMM's device. +FusedAL2Prefetch makeL2Prefetch(torch::TensorList tensors, torch::Device const& device) +{ + TORCH_CHECK(tensors.size() <= static_cast(FusedAL2Prefetch::kMaxRanges), + "glm_fused_a_gemm prefetches at most ", FusedAL2Prefetch::kMaxRanges, " tensors, got ", tensors.size()); + FusedAL2Prefetch ranges{}; + for (size_t i = 0; i < tensors.size(); ++i) + { + torch::Tensor const& tensor = tensors[i]; + TORCH_CHECK(tensor.device() == device, "l2_prefetch tensors must be on the GEMM's device"); + TORCH_CHECK(tensor.is_contiguous(), "l2_prefetch tensors must be contiguous"); + TORCH_CHECK(isAligned(tensor), "l2_prefetch tensors must be 16-byte aligned"); + ranges.ptr[i] = tensor.data_ptr(); + ranges.bytes[i] = static_cast(tensor.numel()) * tensor.element_size(); + } + return ranges; +} + +} // namespace + +//! The GLM-5 fused-A GEMM (kv_a_proj_with_mqa) for up to 8 tokens. +//! +//! @param hiddenStates BF16 CUDA tensor [num_tokens, 6144], contiguous, num_tokens in [1, 8]. +//! @param weight BF16 CUDA tensor [out_features, 6144], contiguous: the kv_a_proj_with_mqa weight (2624 rows: q_a, +//! compressed KV, k_pe) or that weight with the DSA indexer's wk and weights_proj rows appended (2784 rows). +//! @param l2Prefetch Up to four contiguous, 16-byte aligned tensors on the same device that the kernel prefetches into +//! L2 before its grid-dependency wait (weights later kernels of the step read); may be empty. +//! @return BF16 tensor [num_tokens, out_features] = hiddenStates @ weight^T, FP32 accumulation. +torch::Tensor glmFusedAGemm( + torch::Tensor const& hiddenStates, torch::Tensor const& weight, torch::TensorList l2Prefetch) +{ + namespace glm = tk::glm_moe_dsa; + TORCH_CHECK(hiddenStates.is_cuda() && weight.is_cuda(), "glm_fused_a_gemm expects CUDA tensors"); + TORCH_CHECK(hiddenStates.device() == weight.device(), "hidden_states and weight must be on the same device"); + TORCH_CHECK(hiddenStates.scalar_type() == torch::kBFloat16 && weight.scalar_type() == torch::kBFloat16, + "glm_fused_a_gemm expects BF16 hidden_states and weight, got ", hiddenStates.scalar_type(), " and ", + weight.scalar_type()); + TORCH_CHECK(hiddenStates.dim() == 2 && weight.dim() == 2, + "glm_fused_a_gemm expects hidden_states [num_tokens, 6144] and weight [out_features, 6144]"); + TORCH_CHECK(hiddenStates.size(1) == glm::kFusedAHiddenSize && weight.size(1) == glm::kFusedAHiddenSize, + "glm_fused_a_gemm serves hidden size ", glm::kFusedAHiddenSize, ", got hidden_states ", hiddenStates.sizes(), + " and weight ", weight.sizes()); + TORCH_CHECK(weight.size(0) == glm::kFusedAOutFeatures || weight.size(0) == glm::kFusedAIndexerOutFeatures, + "glm_fused_a_gemm serves ", glm::kFusedAOutFeatures, " or ", glm::kFusedAIndexerOutFeatures, + " output features, got ", weight.size(0)); + TORCH_CHECK(hiddenStates.size(0) >= 1 && hiddenStates.size(0) <= glm::kFusedAMaxTokens, + "glm_fused_a_gemm takes 1 to ", glm::kFusedAMaxTokens, " tokens, got ", hiddenStates.size(0)); + TORCH_CHECK(hiddenStates.is_contiguous() && weight.is_contiguous(), + "glm_fused_a_gemm expects contiguous row-major hidden_states and weight"); + TORCH_CHECK(isAligned(hiddenStates) && isAligned(weight), + "glm_fused_a_gemm expects 16-byte aligned hidden_states and weight (tensor copies)"); + + c10::cuda::CUDAGuard const deviceGuard{hiddenStates.device()}; + TORCH_CHECK(tensorrt_llm::common::getSMVersion() >= 90, "glm_fused_a_gemm needs SM90 or later"); + int const numTokens = static_cast(hiddenStates.size(0)); + int const outFeatures = static_cast(weight.size(0)); + torch::Tensor output = torch::empty({numTokens, outFeatures}, hiddenStates.options()); + auto const stream = at::cuda::getCurrentCUDAStream(hiddenStates.get_device()); + auto* outputPtr = reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()); + auto const* hiddenStatesPtr = reinterpret_cast<__nv_bfloat16 const*>(hiddenStates.data_ptr()); + auto const* weightPtr = reinterpret_cast<__nv_bfloat16 const*>(weight.data_ptr()); + if (l2Prefetch.empty()) + { + glm::invokeFusedAGemm(outputPtr, hiddenStatesPtr, weightPtr, numTokens, outFeatures, nullptr, stream); + } + else + { + FusedAL2Prefetch const ranges = makeL2Prefetch(l2Prefetch, hiddenStates.device()); + glm::invokeFusedAGemm(outputPtr, hiddenStatesPtr, weightPtr, numTokens, outFeatures, &ranges, stream); + } + return output; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("glm_fused_a_gemm(Tensor hidden_states, Tensor weight, Tensor[] l2_prefetch) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_fused_a_gemm", &tensorrt_llm::torch_ext::glmFusedAGemm); +} diff --git a/cpp/tensorrt_llm/thop/glmFusedEmbeddingOp.cpp b/cpp/tensorrt_llm/thop/glmFusedEmbeddingOp.cpp new file mode 100644 index 000000000000..cd95ab8210e7 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmFusedEmbeddingOp.cpp @@ -0,0 +1,74 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +// Lookup of an unsharded BF16 embedding table in one kernel. Ids outside [0, weight.size(0)) give a zero row, +// the same guard as tensorrt_llm._torch.modules.embedding.pre_comm_embedding_ops. +at::Tensor glmFusedEmbedding(at::Tensor const& input, at::Tensor const& weight) +{ + CHECK_TH_CUDA(input); + CHECK_TH_CUDA(weight); + TORCH_CHECK(input.device() == weight.device(), "glm_fused_embedding: input and weight must be on the same device"); + TORCH_CHECK(input.scalar_type() == at::ScalarType::Int || input.scalar_type() == at::ScalarType::Long, + "glm_fused_embedding expects int32 or int64 ids, got ", input.scalar_type()); + CHECK_TYPE(weight, at::ScalarType::BFloat16); + CHECK_CONTIGUOUS(weight); + TORCH_CHECK(weight.dim() == 2, "glm_fused_embedding expects a 2D weight, got ", weight.dim(), "D"); + + int64_t const vocabSize = weight.size(0); + int64_t const hiddenSize = weight.size(1); + TORCH_CHECK(hiddenSize % 8 == 0, "glm_fused_embedding needs a hidden size divisible by 8, got ", hiddenSize); + TORCH_CHECK( + reinterpret_cast(weight.data_ptr()) % 16 == 0, "glm_fused_embedding needs a 16-byte aligned weight"); + + at::Tensor const ids = input.contiguous(); + std::vector outShape(input.sizes().begin(), input.sizes().end()); + outShape.push_back(hiddenSize); + at::Tensor const out = at::empty(outShape, weight.options()); + + auto const stream = at::cuda::getCurrentCUDAStream(input.get_device()); + tensorrt_llm::kernels::glm_moe_dsa::invokeFusedEmbeddingLookup(static_cast<__nv_bfloat16*>(out.data_ptr()), + ids.data_ptr(), ids.scalar_type() == at::ScalarType::Long, static_cast<__nv_bfloat16 const*>(weight.data_ptr()), + ids.numel(), vocabSize, hiddenSize, stream); + return out; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("glm_fused_embedding(Tensor input, Tensor weight) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_fused_embedding", &tensorrt_llm::torch_ext::glmFusedEmbedding); +} diff --git a/cpp/tensorrt_llm/thop/glmLowMBf16MoeOp.cpp b/cpp/tensorrt_llm/thop/glmLowMBf16MoeOp.cpp new file mode 100644 index 000000000000..e414e334d7b8 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmLowMBf16MoeOp.cpp @@ -0,0 +1,123 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ +bool isAligned16(at::Tensor const& t) +{ + return reinterpret_cast(t.data_ptr()) % 16 == 0; +} +} // namespace + +// One token through an MoE layer with unquantized BF16 SwiGLU experts held as TP slices and one-group DeepSeek-V3 +// routing: routing from the router logits, FC1 + SwiGLU, and FC2 + the weighted sum over the routed experts, in two +// kernels (glmLowMBf16Moe.h). The weights are read in the layout CutlassFusedMoE keeps: w3_w1 [E, 2I, H] with the up +// (W3) rows first, w2 [E, H, I]. Returns [1, H]. +at::Tensor glmLowmBf16Moe(at::Tensor const& x, at::Tensor const& routerLogits, at::Tensor const& eScoreCorrectionBias, + at::Tensor const& w3W1, at::Tensor const& w2, int64_t topK, double routedScalingFactor) +{ + namespace tk = tensorrt_llm::kernels::glm_moe_dsa; + + CHECK_INPUT(x, at::ScalarType::BFloat16); + CHECK_INPUT(routerLogits, at::ScalarType::Float); + CHECK_TH_CUDA(eScoreCorrectionBias); + CHECK_CONTIGUOUS(eScoreCorrectionBias); + CHECK_INPUT(w3W1, at::ScalarType::BFloat16); + CHECK_INPUT(w2, at::ScalarType::BFloat16); + TORCH_CHECK(eScoreCorrectionBias.scalar_type() == at::ScalarType::Float + || eScoreCorrectionBias.scalar_type() == at::ScalarType::BFloat16, + "glm_lowm_bf16_moe expects an FP32 or BF16 e_score_correction_bias, got ", eScoreCorrectionBias.scalar_type()); + for (at::Tensor const* t : {&routerLogits, &eScoreCorrectionBias, &w3W1, &w2}) + { + TORCH_CHECK(t->device() == x.device(), "glm_lowm_bf16_moe: all tensors must be on the device of x"); + } + TORCH_CHECK(tensorrt_llm::common::getSMVersion() >= 90, "glm_lowm_bf16_moe needs SM90 or newer"); + + TORCH_CHECK(x.dim() == 2 && x.size(0) == 1, "glm_lowm_bf16_moe takes one token, got x of shape ", x.sizes()); + TORCH_CHECK(w3W1.dim() == 3 && w2.dim() == 3, "glm_lowm_bf16_moe expects 3D expert weights"); + int64_t const hiddenSize = x.size(1); + int64_t const numExperts = w3W1.size(0); + int64_t const interSize = w2.size(2); + TORCH_CHECK(numExperts == tk::kLowMBf16MoeNumExperts, "glm_lowm_bf16_moe is built for ", tk::kLowMBf16MoeNumExperts, + " experts, got ", numExperts); + TORCH_CHECK( + topK == tk::kLowMBf16MoeTopK, "glm_lowm_bf16_moe is built for top-", tk::kLowMBf16MoeTopK, ", got top-", topK); + TORCH_CHECK(w3W1.size(1) == 2 * interSize && w3W1.size(2) == hiddenSize, "glm_lowm_bf16_moe: w3_w1 of shape ", + w3W1.sizes(), " does not match x ", x.sizes(), " and w2 ", w2.sizes()); + TORCH_CHECK(w2.size(0) == numExperts && w2.size(1) == hiddenSize, "glm_lowm_bf16_moe: w2 of shape ", w2.sizes(), + " does not match x ", x.sizes(), " and w3_w1 ", w3W1.sizes()); + TORCH_CHECK(routerLogits.dim() == 2 && routerLogits.size(0) == 1 && routerLogits.size(1) == numExperts, + "glm_lowm_bf16_moe: router_logits must be [1, ", numExperts, "], got ", routerLogits.sizes()); + TORCH_CHECK(eScoreCorrectionBias.dim() == 1 && eScoreCorrectionBias.size(0) == numExperts, + "glm_lowm_bf16_moe: e_score_correction_bias must be [", numExperts, "], got ", eScoreCorrectionBias.sizes()); + TORCH_CHECK(tk::lowMBf16MoeSupportsShape(hiddenSize, interSize), "glm_lowm_bf16_moe: unsupported hidden size ", + hiddenSize, " or intermediate size ", interSize); + TORCH_CHECK(isAligned16(x) && isAligned16(w3W1) && isAligned16(w2), + "glm_lowm_bf16_moe needs 16-byte aligned x, w3_w1 and w2"); + + at::Tensor out = at::empty({1, hiddenSize}, x.options()); + at::Tensor act = at::empty({topK, interSize}, x.options()); + auto const stream = at::cuda::getCurrentCUDAStream(x.get_device()); + auto* outPtr = static_cast<__nv_bfloat16*>(out.data_ptr()); + auto* actPtr = static_cast<__nv_bfloat16*>(act.data_ptr()); + auto const* xPtr = static_cast<__nv_bfloat16 const*>(x.data_ptr()); + auto const* logitsPtr = static_cast(routerLogits.data_ptr()); + auto const* w3W1Ptr = static_cast<__nv_bfloat16 const*>(w3W1.data_ptr()); + auto const* w2Ptr = static_cast<__nv_bfloat16 const*>(w2.data_ptr()); + if (eScoreCorrectionBias.scalar_type() == at::ScalarType::Float) + { + tk::invokeLowMBf16Moe(outPtr, actPtr, xPtr, logitsPtr, + static_cast(eScoreCorrectionBias.data_ptr()), w3W1Ptr, w2Ptr, static_cast(hiddenSize), + static_cast(interSize), routedScalingFactor, stream); + } + else + { + tk::invokeLowMBf16Moe<__nv_bfloat16>(outPtr, actPtr, xPtr, logitsPtr, + static_cast<__nv_bfloat16 const*>(eScoreCorrectionBias.data_ptr()), w3W1Ptr, w2Ptr, + static_cast(hiddenSize), static_cast(interSize), routedScalingFactor, stream); + } + return out; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_lowm_bf16_moe(Tensor x, Tensor router_logits, Tensor e_score_correction_bias, Tensor w3_w1, Tensor w2, " + "int top_k, float routed_scaling_factor) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_lowm_bf16_moe", &tensorrt_llm::torch_ext::glmLowmBf16Moe); +} diff --git a/cpp/tensorrt_llm/thop/glmMlaKvHalfAppendOp.cpp b/cpp/tensorrt_llm/thop/glmMlaKvHalfAppendOp.cpp new file mode 100644 index 000000000000..c2d2759a08f4 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmMlaKvHalfAppendOp.cpp @@ -0,0 +1,148 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/quantization.h" +#include "tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.h" +#include "tensorrt_llm/runtime/torchUtils.h" +#include "tensorrt_llm/thop/attentionOp.h" + +#include +#include + +#include +#include +#include + +namespace th = torch; +namespace tc = tensorrt_llm::common; +namespace tk = tensorrt_llm::kernels::glm_moe_dsa; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kInt32Max = std::numeric_limits::max(); +constexpr int64_t kLatentRowDim = tk::kGlmKvHalfLatentDim + tk::kGlmKvHalfRopeDim; +// The kernel walks the latent rows with 16-byte vector loads. +constexpr int64_t kVectorBytes = 16; +// Per position, the rope table holds 64 (cos, sin) float pairs. +constexpr int64_t kRopeTableFloatsPerPosition = 2 * tk::kGlmKvHalfRopeDim; + +} // namespace + +// The KV half of GLM-5.2's MLA generation step: kv_a_layernorm of the 512 latent columns of the raw kv_a_proj output, +// RoPE of its 64 K columns, E4M3 at unit scale, appended to layer layer_idx of the paged FP8 cache at the generation +// positions sequence_length[num_contexts + s] - seq_len_q + t, plus the generation attention's scales in +// mla_bmm1_scale / mla_bmm2_scale. The paged-cache arguments are those of mla_rope_generation. +void glmMlaKvHalfAppend(th::Tensor const& latent, th::Tensor const& kvNormWeight, double kvNormEps, + th::Tensor const& rotaryCosSin, th::Tensor const& sequenceLength, int64_t numContexts, int64_t seqLenQ, + th::Tensor const& kvCacheBlockOffsets, th::Tensor const& hostKvCachePoolPointers, + th::Tensor const& hostKvCachePoolMapping, int64_t layerIdx, int64_t tokensPerBlock, int64_t attentionWindowSize, + th::Tensor& mlaBmm1Scale, th::Tensor& mlaBmm2Scale, double softmaxScale, bool triggerAtEntry) +{ + TORCH_CHECK(latent.is_cuda() && kvNormWeight.is_cuda() && rotaryCosSin.is_cuda() && sequenceLength.is_cuda() + && kvCacheBlockOffsets.is_cuda() && mlaBmm1Scale.is_cuda() && mlaBmm2Scale.is_cuda(), + "glm_mla_kv_half_append takes CUDA tensors besides the host pool pointers and mapping"); + auto const device = latent.device(); + TORCH_CHECK(kvNormWeight.device() == device && rotaryCosSin.device() == device && sequenceLength.device() == device + && kvCacheBlockOffsets.device() == device && mlaBmm1Scale.device() == device + && mlaBmm2Scale.device() == device, + "glm_mla_kv_half_append takes its CUDA tensors on one device"); + + TORCH_CHECK(latent.scalar_type() == th::kBFloat16 && latent.dim() == 2 && latent.size(1) >= kLatentRowDim + && latent.stride(1) == 1, + "latent must be a BF16 [tokens, >= ", kLatentRowDim, "] view with unit column stride"); + TORCH_CHECK(latent.stride(0) % (kVectorBytes / 2) == 0 && latent.stride(0) <= kInt32Max + && reinterpret_cast(latent.data_ptr()) % kVectorBytes == 0, + "latent rows must be 16-byte aligned (row stride a multiple of 8 elements), got row stride ", latent.stride(0)); + int64_t const numTokens = latent.size(0); + TORCH_CHECK(numTokens <= kInt32Max, "too many latent rows"); + TORCH_CHECK(seqLenQ > 0 && numTokens % seqLenQ == 0, "seq_len_q must be positive and divide the latent rows; got ", + seqLenQ, " for ", numTokens, " rows"); + int64_t const numGenerations = numTokens / seqLenQ; + + TORCH_CHECK(kvNormWeight.scalar_type() == th::kBFloat16 && kvNormWeight.is_contiguous() + && kvNormWeight.numel() == tk::kGlmKvHalfLatentDim + && reinterpret_cast(kvNormWeight.data_ptr()) % kVectorBytes == 0, + "kv_norm_weight must be a contiguous, 16-byte aligned BF16 tensor of ", tk::kGlmKvHalfLatentDim, " elements"); + TORCH_CHECK(rotaryCosSin.scalar_type() == th::kFloat32 && rotaryCosSin.is_contiguous() + && rotaryCosSin.numel() % kRopeTableFloatsPerPosition == 0, + "rotary_cos_sin must be a contiguous float32 table of ", kRopeTableFloatsPerPosition, " values per position"); + TORCH_CHECK(sequenceLength.scalar_type() == th::kInt32 && sequenceLength.dim() == 1 + && sequenceLength.is_contiguous() && numContexts >= 0 + && sequenceLength.size(0) >= numContexts + numGenerations, + "sequence_length must be a contiguous int32 tensor covering the context and generation sequences"); + // KVBlockArray addresses sequence s of the pool's rows at s * 2 * max_blocks. + TORCH_CHECK(kvCacheBlockOffsets.scalar_type() == th::kInt32 && kvCacheBlockOffsets.dim() == 4 + && kvCacheBlockOffsets.size(2) == 2 && kvCacheBlockOffsets.stride(3) == 1 + && kvCacheBlockOffsets.stride(2) == kvCacheBlockOffsets.size(3) + && kvCacheBlockOffsets.stride(1) == 2 * kvCacheBlockOffsets.size(3) + && kvCacheBlockOffsets.size(1) >= numContexts + numGenerations, + "kv_cache_block_offsets must be an int32 [pools, sequences, 2, max_blocks] tensor, dense in its last three " + "dimensions"); + TORCH_CHECK(hostKvCachePoolPointers.device().is_cpu() && hostKvCachePoolMapping.device().is_cpu(), + "host_kv_cache_pool_pointers and host_kv_cache_pool_mapping must be CPU tensors"); + TORCH_CHECK(tokensPerBlock > 0 && (tokensPerBlock & (tokensPerBlock - 1)) == 0, + "tokens_per_block must be a power of two, got ", tokensPerBlock); + TORCH_CHECK(attentionWindowSize > 0 && attentionWindowSize <= kInt32Max, "invalid attention_window_size ", + attentionWindowSize); + TORCH_CHECK(mlaBmm1Scale.scalar_type() == th::kFloat32 && mlaBmm1Scale.is_contiguous() && mlaBmm1Scale.numel() >= 2, + "mla_bmm1_scale must hold two contiguous float32 values"); + TORCH_CHECK(mlaBmm2Scale.scalar_type() == th::kFloat32 && mlaBmm2Scale.numel() >= 1, + "mla_bmm2_scale must hold a float32 value"); + + // The latent pool of the FP8 MLA cache: one 576-element head per token, one byte per element. + auto const kvCacheBuffers = buildPagedKvCacheBuffers(std::optional(kvCacheBlockOffsets), + std::optional(hostKvCachePoolPointers), std::optional(hostKvCachePoolMapping), + tc::QuantMode::fp8KvCache(), layerIdx, numGenerations, tokensPerBlock, /*kv_head_num=*/1, kLatentRowDim, + attentionWindowSize, attentionWindowSize, /*beam_width=*/1, numContexts, /*is_mla_enable=*/true, + latent.element_size()); + TORCH_CHECK(kvCacheBuffers.kvCacheBuffer.data != nullptr, "glm_mla_kv_half_append found no paged KV cache"); + + auto const stream = at::cuda::getCurrentCUDAStream(latent.get_device()); + tk::invokeGlmMlaKvHalfAppend(static_cast<__nv_bfloat16 const*>(latent.data_ptr()), + static_cast(latent.stride(0)), kvCacheBuffers.kvCacheBuffer, + reinterpret_cast(rotaryCosSin.data_ptr()), + static_cast<__nv_bfloat16 const*>(kvNormWeight.data_ptr()), static_cast(kvNormEps), + static_cast(numTokens), static_cast(seqLenQ), + sequenceLength.data_ptr() + numContexts, mlaBmm1Scale.data_ptr(), + mlaBmm2Scale.data_ptr(), static_cast(softmaxScale), triggerAtEntry, stream); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_mla_kv_half_append(Tensor latent, Tensor kv_norm_weight, float kv_norm_eps, Tensor rotary_cos_sin, " + "Tensor sequence_length, int num_contexts, int seq_len_q, Tensor kv_cache_block_offsets, " + "Tensor host_kv_cache_pool_pointers, Tensor host_kv_cache_pool_mapping, int layer_idx, int tokens_per_block, " + "int attention_window_size, Tensor(a!) mla_bmm1_scale, Tensor(b!) mla_bmm2_scale, float softmax_scale, " + "bool trigger_at_entry=False) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_mla_kv_half_append", &tensorrt_llm::torch_ext::glmMlaKvHalfAppend); +} diff --git a/cpp/tensorrt_llm/thop/glmMnnvlAllgatherOp.cpp b/cpp/tensorrt_llm/thop/glmMnnvlAllgatherOp.cpp new file mode 100644 index 000000000000..ff992084c3d2 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmMnnvlAllgatherOp.cpp @@ -0,0 +1,100 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/mcastDevMemUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.h" +#include "tensorrt_llm/runtime/mcastDeviceMemory.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +// All-gather of BF16 rows along their last dimension, widened to FP32, over a dedicated MNNVL exchange workspace +// (comm_buffer, the unicast view of its three Lamport buffers, and buffer_flags): the rank-ordered concatenation of +// the shards (kernels/glmMoeDsa/glmMnnvlAllgather.h). +at::Tensor glmMnnvlAllgather( + at::Tensor const& input, at::Tensor& comm_buffer, at::Tensor& buffer_flags, int64_t group_size) +{ + char const* op = "glm_mnnvl_allgather"; + auto* mcastMem = tensorrt_llm::common::findMcastDevMemBuffer(comm_buffer.data_ptr()); + TORCH_CHECK(mcastMem != nullptr, op, ": comm_buffer must be obtained from a McastGPUBuffer instance"); + TORCH_CHECK(mcastMem->isMapped(), op, ": the MNNVL workspace handles are not attached"); + CHECK_TH_CUDA(input); + CHECK_CONTIGUOUS(input); + TORCH_CHECK(input.dim() >= 1, op, ": input must have at least one dimension"); + TORCH_CHECK(input.scalar_type() == at::ScalarType::BFloat16, op, ": input must be BF16, got ", input.scalar_type()); + TORCH_CHECK(static_cast(mcastMem->getWorldSize()) == group_size, op, ": group_size ", group_size, + " does not match the workspace group of ", mcastMem->getWorldSize(), " ranks"); + TORCH_CHECK(buffer_flags.element_size() == 4 && buffer_flags.numel() >= 9 && buffer_flags.is_contiguous(), op, + ": buffer_flags must hold the workspace's nine 32-bit Lamport flags"); + TORCH_CHECK(comm_buffer.device() == input.device() && buffer_flags.device() == input.device(), op, + ": the workspace must be on the device of input"); + + int64_t const shardDim = input.size(-1); + int64_t const numRows = shardDim == 0 ? 0 : input.numel() / shardDim; + TORCH_CHECK(shardDim * input.element_size() % 16 == 0, op, ": a shard row must be a multiple of 16 bytes"); + TORCH_CHECK(reinterpret_cast(input.data_ptr()) % 16 == 0, op, ": input must be 16-byte aligned"); + int64_t const gatheredBytes = numRows * group_size * shardDim * input.element_size(); + int64_t const bufferBytes = comm_buffer.numel() * comm_buffer.element_size() / 3; + TORCH_CHECK(gatheredBytes <= bufferBytes, op, ": ", gatheredBytes, " gathered bytes exceed the Lamport buffer of ", + bufferBytes, " bytes"); + + std::vector outShape(input.sizes().begin(), input.sizes().end()); + outShape.back() = group_size * shardDim; + at::Tensor out = at::empty(outShape, input.options().dtype(at::ScalarType::Float)); + if (numRows == 0) + { + return out; + } + + tensorrt_llm::kernels::glm_moe_dsa::MnnvlAllgatherParams params{}; + params.nRanks = static_cast(group_size); + params.rank = static_cast(mcastMem->getRank()); + params.numRows = static_cast(numRows); + params.shardDim = static_cast(shardDim); + params.input = input.const_data_ptr(); + params.output = out.data_ptr(); + params.bufferPtrLocal = comm_buffer.data_ptr(); + params.multicastPtr = mcastMem->getMulticastPtr(); + params.bufferFlags = reinterpret_cast(buffer_flags.data_ptr()); + params.stream = at::cuda::getCurrentCUDAStream(input.get_device()); + tensorrt_llm::kernels::glm_moe_dsa::mnnvlAllgather(params); + return out; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_mnnvl_allgather(Tensor input, Tensor(a!) comm_buffer, Tensor(b!) buffer_flags, int group_size) -> " + "Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_mnnvl_allgather", &tensorrt_llm::torch_ext::glmMnnvlAllgather); +} diff --git a/cpp/tensorrt_llm/thop/glmMnnvlHopOp.cpp b/cpp/tensorrt_llm/thop/glmMnnvlHopOp.cpp new file mode 100644 index 000000000000..2504d971622d --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmMnnvlHopOp.cpp @@ -0,0 +1,217 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/mcastDevMemUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.h" +#include "tensorrt_llm/runtime/mcastDeviceMemory.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ +namespace +{ + +namespace glm = tensorrt_llm::kernels::glm_moe_dsa; + +bool isAlignedBf16(at::Tensor const& t) +{ + return t.scalar_type() == at::ScalarType::BFloat16 && t.is_contiguous() + && reinterpret_cast(t.data_ptr()) % 16 == 0; +} + +bool hopShapeSupported(int64_t numTokens, int64_t hiddenSize, int64_t nRanks) +{ + constexpr int64_t kIntMax = std::numeric_limits::max(); + if (numTokens > kIntMax || hiddenSize > kIntMax || nRanks > kIntMax) + { + return false; + } + return glm::mnnvlHopSupported(static_cast(numTokens), static_cast(hiddenSize), static_cast(nRanks)); +} + +void checkRowsAndGamma(at::Tensor const& residual, at::Tensor const& gamma, char const* op) +{ + TORCH_CHECK(residual.is_cuda() && residual.dim() == 2 && isAlignedBf16(residual), op, + ": residual must be a contiguous, 16-byte aligned BF16 [num_tokens, hidden] CUDA tensor"); + TORCH_CHECK(gamma.dim() == 1 && gamma.size(0) == residual.size(1) && isAlignedBf16(gamma), op, + ": gamma must be a contiguous, 16-byte aligned BF16 [hidden] tensor"); + TORCH_CHECK(gamma.device() == residual.device(), op, ": gamma must be on the device of the rows"); +} + +// The MNNVL all-reduce workspace that comm_buffer views ([3, N]: this rank's three Lamport buffers) for a hop of the +// rows of residual: the shape must be one the hops serve and its Lamport data must fit one buffer. +tensorrt_llm::runtime::McastDeviceMemory* hopWorkspace( + at::Tensor const& comm_buffer, at::Tensor const& buffer_flags, at::Tensor const& residual, char const* op) +{ + auto* mcastMem = tensorrt_llm::common::findMcastDevMemBuffer(comm_buffer.data_ptr()); + TORCH_CHECK(mcastMem != nullptr, op, ": comm_buffer must be obtained from a McastGPUBuffer instance"); + TORCH_CHECK(mcastMem->isMapped(), op, ": the MNNVL workspace handles are not attached"); + TORCH_CHECK(comm_buffer.dim() == 2 && comm_buffer.size(0) == 3 && comm_buffer.is_contiguous(), op, + ": comm_buffer must be the [3, N] view of the workspace's three Lamport buffers"); + TORCH_CHECK(buffer_flags.element_size() == 4 && buffer_flags.numel() >= 9 && buffer_flags.is_contiguous(), op, + ": buffer_flags must hold the workspace's nine 32-bit Lamport flags"); + TORCH_CHECK(comm_buffer.device() == residual.device() && buffer_flags.device() == residual.device(), op, + ": the workspace must be on the device of the rows"); + int64_t const numTokens = residual.size(0); + int64_t const hiddenSize = residual.size(1); + int64_t const nRanks = static_cast(mcastMem->getWorldSize()); + TORCH_CHECK(hopShapeSupported(numTokens, hiddenSize, nRanks), op, ": unsupported shape: ", numTokens, + " tokens, hidden ", hiddenSize, ", ", nRanks, " ranks"); + int64_t const bufferBytes = comm_buffer.size(1) * comm_buffer.element_size(); + int64_t const hopBytes + = glm::mnnvlHopBufferBytes(static_cast(numTokens), static_cast(hiddenSize), static_cast(nRanks)); + TORCH_CHECK(hopBytes <= bufferBytes, op, ": ", hopBytes, " bytes of Lamport data exceed the Lamport buffer of ", + bufferBytes, " bytes"); + return mcastMem; +} + +glm::MnnvlHopParams makeHopParams(tensorrt_llm::runtime::McastDeviceMemory* mcastMem, at::Tensor& buffer_flags, + at::Tensor const& residual, at::Tensor const& gamma, double epsilon, at::Tensor& normOut, at::Tensor& residualOut) +{ + glm::MnnvlHopParams params{}; + params.nRanks = static_cast(mcastMem->getWorldSize()); + params.rank = static_cast(mcastMem->getRank()); + params.numTokens = static_cast(residual.size(0)); + params.tokenDim = static_cast(residual.size(1)); + params.bufferPtrsDev = mcastMem->getBufferPtrsDev(); + params.multicastPtr = mcastMem->getMulticastPtr(); + params.bufferFlags = reinterpret_cast(buffer_flags.mutable_data_ptr()); + params.input = nullptr; + params.residualIn = residual.const_data_ptr(); + params.gamma = gamma.const_data_ptr(); + params.epsilon = static_cast(epsilon); + params.residualOut = residualOut.mutable_data_ptr(); + params.normOut = normOut.mutable_data_ptr(); + params.shardPublished = false; + params.stream = at::cuda::getCurrentCUDAStream(residual.get_device()); + return params; +} + +} // namespace + +// The one-shot MNNVL all-reduce of BF16 rows with the residual add and RMSNorm (kernels/glmMoeDsa/glmMnnvlHop.h) on +// the MNNVL all-reduce workspace comm_buffer / buffer_flags. Returns [norm_out, residual_out]. With shard_published the +// producer of input already stored this rank's rows into the workspace (the o_proj GEMM's Lamport publish), and input +// only gives the shape. +std::vector glmMnnvlHop(at::Tensor const& input, at::Tensor const& residual, at::Tensor const& gamma, + double epsilon, at::Tensor& comm_buffer, at::Tensor& buffer_flags, bool shard_published) +{ + char const* op = "glm_mnnvl_hop"; + checkRowsAndGamma(residual, gamma, op); + TORCH_CHECK(input.sizes() == residual.sizes(), op, ": input must have the shape of residual"); + if (!shard_published) + { + TORCH_CHECK(isAlignedBf16(input) && input.device() == residual.device(), op, + ": input must be a contiguous, 16-byte aligned BF16 tensor on the device of residual"); + } + at::Tensor normOut = at::empty_like(residual); + at::Tensor residualOut = at::empty_like(residual); + if (residual.size(0) == 0) + { + return {normOut, residualOut}; + } + auto* mcastMem = hopWorkspace(comm_buffer, buffer_flags, residual, op); + auto params = makeHopParams(mcastMem, buffer_flags, residual, gamma, epsilon, normOut, residualOut); + params.input = shard_published ? nullptr : input.const_data_ptr(); + params.shardPublished = shard_published; + glm::mnnvlHop(params); + return {normOut, residualOut}; +} + +// The TRTLLM-Gen MoE finalize (top-8, BF16) fused into the hop: the inputs are the MoE's do_finalize=false outputs (FC2 +// rows, routing weights, expanded-to-permuted index) and the addend (the shared experts' output). Returns [norm_out, +// residual_out], those of glm_mnnvl_hop on the finalized rows. The routing weights and index are read before the +// kernel's grid dependency wait: the kernel that wrote them must have completed when this one launches. +std::vector glmMnnvlMoeFinalizeHop(at::Tensor const& gemm2_output, at::Tensor const& expert_weights, + at::Tensor const& expanded_idx_to_permuted_idx, at::Tensor const& addend, at::Tensor const& residual, + at::Tensor const& gamma, double epsilon, at::Tensor& comm_buffer, at::Tensor& buffer_flags) +{ + char const* op = "glm_mnnvl_moe_finalize_hop"; + checkRowsAndGamma(residual, gamma, op); + int64_t const numTokens = residual.size(0); + TORCH_CHECK(gemm2_output.dim() == 2 && isAlignedBf16(gemm2_output), op, + ": gemm2_output must be a contiguous, 16-byte aligned BF16 2D tensor"); + TORCH_CHECK(gemm2_output.size(1) >= residual.size(1) && gemm2_output.size(1) % 8 == 0, op, + ": gemm2_output rows must hold the hidden dimension in a multiple of 8 elements"); + TORCH_CHECK(expanded_idx_to_permuted_idx.scalar_type() == at::ScalarType::Int + && expanded_idx_to_permuted_idx.dim() == 2 && expanded_idx_to_permuted_idx.is_contiguous() + && expanded_idx_to_permuted_idx.size(0) == numTokens + && expanded_idx_to_permuted_idx.size(1) == glm::kMnnvlHopMoeTopK, + op, ": expanded_idx_to_permuted_idx must be contiguous int32 [num_tokens, 8]"); + TORCH_CHECK(expert_weights.scalar_type() == at::ScalarType::BFloat16 && expert_weights.is_contiguous() + && expert_weights.sizes() == expanded_idx_to_permuted_idx.sizes(), + op, ": expert_weights must be contiguous BF16 of the index's shape"); + TORCH_CHECK(addend.sizes() == residual.sizes() && isAlignedBf16(addend), op, + ": addend must be a contiguous, 16-byte aligned BF16 tensor of the shape of residual"); + for (auto const* t : {&gemm2_output, &expert_weights, &expanded_idx_to_permuted_idx, &addend}) + { + TORCH_CHECK(t->device() == residual.device(), op, ": all inputs must be on one device"); + } + at::Tensor normOut = at::empty_like(residual); + at::Tensor residualOut = at::empty_like(residual); + if (numTokens == 0) + { + return {normOut, residualOut}; + } + auto* mcastMem = hopWorkspace(comm_buffer, buffer_flags, residual, op); + auto const params = makeHopParams(mcastMem, buffer_flags, residual, gamma, epsilon, normOut, residualOut); + glm::MnnvlMoeFinalizeParams moe{}; + moe.fc2Output = gemm2_output.const_data_ptr(); + moe.hiddenDimPadded = static_cast(gemm2_output.size(1)); + moe.expertWeights = expert_weights.const_data_ptr(); + moe.expandedIdxToPermutedIdx = expanded_idx_to_permuted_idx.const_data_ptr(); + moe.addend = addend.const_data_ptr(); + glm::mnnvlMoeFinalizeHop(params, moe); + return {normOut, residualOut}; +} + +// Whether glm_mnnvl_hop and glm_mnnvl_moe_finalize_hop serve [num_tokens, hidden_size] BF16 rows over num_ranks ranks +// on the current device. +bool glmMnnvlHopSupported(int64_t num_tokens, int64_t hidden_size, int64_t num_ranks) +{ + return hopShapeSupported(num_tokens, hidden_size, num_ranks); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_mnnvl_hop(Tensor input, Tensor residual, Tensor gamma, float epsilon, Tensor(a!) comm_buffer, " + "Tensor(b!) buffer_flags, bool shard_published) -> Tensor[]"); + m.def( + "glm_mnnvl_moe_finalize_hop(Tensor gemm2_output, Tensor expert_weights, Tensor expanded_idx_to_permuted_idx, " + "Tensor addend, Tensor residual, Tensor gamma, float epsilon, Tensor(a!) comm_buffer, Tensor(b!) buffer_flags) " + "-> Tensor[]"); + m.def("glm_mnnvl_hop_supported", &tensorrt_llm::torch_ext::glmMnnvlHopSupported); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_mnnvl_hop", &tensorrt_llm::torch_ext::glmMnnvlHop); + m.impl("glm_mnnvl_moe_finalize_hop", &tensorrt_llm::torch_ext::glmMnnvlMoeFinalizeHop); +} diff --git a/cpp/tensorrt_llm/thop/glmRouterGemmOp.cpp b/cpp/tensorrt_llm/thop/glmRouterGemmOp.cpp new file mode 100644 index 000000000000..cbaa789904ce --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmRouterGemmOp.cpp @@ -0,0 +1,94 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.h" + +#include +#include +#include +#include + +namespace tk = tensorrt_llm::kernels; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +// The kernel reads both operands in 16-byte vectors. +constexpr int64_t kVectorBytes = 16; + +bool isAligned(torch::Tensor const& tensor) +{ + return reinterpret_cast(tensor.data_ptr()) % kVectorBytes == 0; +} + +} // namespace + +//! The GLM-5 MoE router logits for up to 16 tokens. +//! +//! @param hiddenStates BF16 CUDA tensor [num_tokens, 6144], contiguous, num_tokens in [1, 16]. +//! @param weight BF16 CUDA tensor [256, 6144], contiguous: the router (gate) weight as stored. +//! @return FP32 tensor [num_tokens, 256] = hiddenStates @ weight^T, FP32 accumulation. +torch::Tensor glmRouterGemm(torch::Tensor const& hiddenStates, torch::Tensor const& weight) +{ + namespace glm = tk::glm_moe_dsa; + TORCH_CHECK(hiddenStates.is_cuda() && weight.is_cuda(), "glm_router_gemm expects CUDA tensors"); + TORCH_CHECK(hiddenStates.device() == weight.device(), "hidden_states and weight must be on the same device"); + TORCH_CHECK(hiddenStates.scalar_type() == torch::kBFloat16 && weight.scalar_type() == torch::kBFloat16, + "glm_router_gemm expects BF16 hidden_states and weight, got ", hiddenStates.scalar_type(), " and ", + weight.scalar_type()); + TORCH_CHECK(hiddenStates.dim() == 2 && weight.dim() == 2, + "glm_router_gemm expects hidden_states [num_tokens, 6144] and weight [256, 6144]"); + TORCH_CHECK(hiddenStates.size(1) == glm::kRouterHiddenSize && weight.size(1) == glm::kRouterHiddenSize, + "glm_router_gemm serves hidden size ", glm::kRouterHiddenSize, ", got hidden_states ", hiddenStates.sizes(), + " and weight ", weight.sizes()); + TORCH_CHECK(weight.size(0) == glm::kRouterNumExperts, "glm_router_gemm serves ", glm::kRouterNumExperts, + " experts, got ", weight.size(0)); + TORCH_CHECK(hiddenStates.size(0) >= 1 && hiddenStates.size(0) <= glm::kRouterMaxTokens, + "glm_router_gemm takes 1 to ", glm::kRouterMaxTokens, " tokens, got ", hiddenStates.size(0)); + TORCH_CHECK(hiddenStates.is_contiguous() && weight.is_contiguous(), + "glm_router_gemm expects contiguous row-major hidden_states and weight"); + TORCH_CHECK(isAligned(hiddenStates) && isAligned(weight), + "glm_router_gemm expects 16-byte aligned hidden_states and weight"); + + c10::cuda::CUDAGuard const deviceGuard{hiddenStates.device()}; + int const numTokens = static_cast(hiddenStates.size(0)); + torch::Tensor logits + = torch::empty({numTokens, glm::kRouterNumExperts}, hiddenStates.options().dtype(torch::kFloat32)); + auto const stream = at::cuda::getCurrentCUDAStream(hiddenStates.get_device()); + glm::invokeRouterGemm(logits.data_ptr(), reinterpret_cast<__nv_bfloat16 const*>(hiddenStates.data_ptr()), + reinterpret_cast<__nv_bfloat16 const*>(weight.data_ptr()), numTokens, stream); + return logits; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("glm_router_gemm(Tensor hidden_states, Tensor weight) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_router_gemm", &tensorrt_llm::torch_ext::glmRouterGemm); +} diff --git a/cpp/tensorrt_llm/thop/glmSparseMlaFlashDecodeOp.cpp b/cpp/tensorrt_llm/thop/glmSparseMlaFlashDecodeOp.cpp new file mode 100644 index 000000000000..019d75820aa9 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmSparseMlaFlashDecodeOp.cpp @@ -0,0 +1,178 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.h" +#include "tensorrt_llm/runtime/torchUtils.h" + +#include + +#include +#include +#include + +namespace th = torch; +namespace tk = tensorrt_llm::kernels::glm_moe_dsa; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kInt32Max = std::numeric_limits::max(); +// The grid spans the query tokens in its y dimension. +constexpr int64_t kMaxQueryTokens = 65535; +constexpr int64_t kBulkCopyAlignment = 16; + +bool isE4m3Storage(th::Tensor const& t) +{ + return t.scalar_type() == th::kFloat8_e4m3fn || t.scalar_type() == th::kByte; +} + +bool isBulkCopyAligned(th::Tensor const& t) +{ + return reinterpret_cast(t.data_ptr()) % kBulkCopyAlignment == 0; +} + +} // namespace + +// Generation sparse MLA of GLM-5.2 at 4 query heads per rank, with the v_b BMM folded in: writes +// output[t][h * 256 + v] = sum_d bf16(attention(q[t, h]))[d] * v_b_proj[h][v][d] for the first q.size(0) rows of +// output. With block_table, indices hold token positions that the kernel maps to pool rows as +// convert_req_index_to_global does (block_table[s][p / tokens_per_block] * page_index_scale * tokens_per_block +// + p % tokens_per_block + layer_offset * tokens_per_block); without it, indices hold pool rows. +void glmSparseMlaFlashDecode(th::Tensor const& q, th::Tensor const& kvPool, th::Tensor const& indices, + th::Tensor const& seqLensKv, th::Tensor const& mlaBmm1Scale, th::Tensor const& mlaBmm2Scale, + th::Tensor const& vBProj, th::Tensor& output, int64_t seqLenQ, int64_t topK, + std::optional const& blockTable, int64_t tokensPerBlock, int64_t pageIndexScale, int64_t layerOffset, + bool gatherBeforeWait) +{ + TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "glm_sparse_mla_flash_decode needs an SM100-family GPU"); + TORCH_CHECK(q.is_cuda() && kvPool.is_cuda() && indices.is_cuda() && seqLensKv.is_cuda() && mlaBmm1Scale.is_cuda() + && mlaBmm2Scale.is_cuda() && vBProj.is_cuda() && output.is_cuda(), + "glm_sparse_mla_flash_decode takes CUDA tensors"); + auto const device = q.device(); + TORCH_CHECK(kvPool.device() == device && indices.device() == device && seqLensKv.device() == device + && mlaBmm1Scale.device() == device && mlaBmm2Scale.device() == device && vBProj.device() == device + && output.device() == device, + "glm_sparse_mla_flash_decode takes tensors on one device"); + + TORCH_CHECK(isE4m3Storage(q), "q must hold E4M3 values (float8_e4m3fn or uint8)"); + TORCH_CHECK(q.dim() == 3 && q.size(1) == tk::kGlmSparseMlaNumHeads && q.size(2) == tk::kGlmSparseMlaHeadDimQk + && q.is_contiguous() && isBulkCopyAligned(q), + "q must be a contiguous, 16-byte aligned [tokens, ", tk::kGlmSparseMlaNumHeads, ", ", + tk::kGlmSparseMlaHeadDimQk, "] tensor"); + int64_t const numTokens = q.size(0); + TORCH_CHECK(numTokens <= kMaxQueryTokens, "glm_sparse_mla_flash_decode takes at most ", kMaxQueryTokens, + " query tokens, got ", numTokens); + TORCH_CHECK(seqLenQ > 0 && seqLenQ <= kInt32Max && numTokens % seqLenQ == 0, + "seq_len_q must be positive and divide the query tokens; got ", seqLenQ, " for ", numTokens, " tokens"); + int64_t const numSeqs = numTokens / seqLenQ; + + TORCH_CHECK(isE4m3Storage(kvPool), "kv_pool must hold E4M3 values (float8_e4m3fn or uint8)"); + TORCH_CHECK(kvPool.is_contiguous() && isBulkCopyAligned(kvPool) && kvPool.numel() % tk::kGlmSparseMlaHeadDimQk == 0, + "kv_pool must be a contiguous, 16-byte aligned pool of ", tk::kGlmSparseMlaHeadDimQk, "-byte token rows"); + TORCH_CHECK(kvPool.numel() / tk::kGlmSparseMlaHeadDimQk <= kInt32Max, "kv_pool rows must be int32-addressable"); + + TORCH_CHECK(topK > 0 && topK <= tk::kGlmSparseMlaMaxTopK, "top_k must be in [1, ", tk::kGlmSparseMlaMaxTopK, + "], got ", topK); + TORCH_CHECK(indices.scalar_type() == th::kInt32 && indices.dim() == 2 && indices.size(0) == numTokens + && indices.size(1) >= topK && indices.stride(1) == 1 && indices.stride(0) >= topK + && indices.stride(0) <= kInt32Max, + "indices must be an int32 [tokens, >= top_k] tensor with unit column stride"); + TORCH_CHECK(seqLensKv.scalar_type() == th::kInt32 && seqLensKv.dim() == 1 && seqLensKv.size(0) >= numSeqs + && seqLensKv.is_contiguous(), + "seq_lens_kv must be a contiguous int32 tensor with one entry per sequence"); + TORCH_CHECK(mlaBmm1Scale.scalar_type() == th::kFloat32 && mlaBmm1Scale.is_contiguous() && mlaBmm1Scale.numel() >= 2, + "mla_bmm1_scale must hold two contiguous float32 values"); + TORCH_CHECK(mlaBmm2Scale.scalar_type() == th::kFloat32 && mlaBmm2Scale.numel() >= 1, + "mla_bmm2_scale must hold a float32 value"); + + TORCH_CHECK(vBProj.scalar_type() == th::kBFloat16 && vBProj.dim() == 3 + && vBProj.size(0) == tk::kGlmSparseMlaNumHeads && vBProj.size(1) == tk::kGlmSparseMlaVHeadDim + && vBProj.size(2) == tk::kGlmSparseMlaHeadDimV && vBProj.is_contiguous() && isBulkCopyAligned(vBProj), + "v_b_proj must be a contiguous, 16-byte aligned BF16 [", tk::kGlmSparseMlaNumHeads, ", ", + tk::kGlmSparseMlaVHeadDim, ", ", tk::kGlmSparseMlaHeadDimV, "] tensor"); + TORCH_CHECK(output.scalar_type() == th::kBFloat16 && output.dim() == 2 && output.size(0) >= numTokens + && output.size(1) == tk::kGlmSparseMlaNumHeads * tk::kGlmSparseMlaVHeadDim && output.stride(1) == 1, + "output must be a BF16 [>= tokens, ", tk::kGlmSparseMlaNumHeads * tk::kGlmSparseMlaVHeadDim, + "] tensor with unit column stride"); + + tk::GlmSparseMlaFlashDecodeParams params{}; + params.q = q.data_ptr(); + params.kvPool = kvPool.data_ptr(); + params.indices = indices.data_ptr(); + params.seqLensKv = seqLensKv.data_ptr(); + params.scaleSoftmaxLog2 = mlaBmm1Scale.data_ptr() + 1; + params.outputScale = mlaBmm2Scale.data_ptr(); + params.vbProj = vBProj.data_ptr(); + params.out = output.data_ptr(); + params.outStride = output.stride(0); + params.numTokensQ = static_cast(numTokens); + params.seqLenQ = static_cast(seqLenQ); + params.indicesStride = static_cast(indices.stride(0)); + params.topK = static_cast(topK); + params.gatherBeforeWait = gatherBeforeWait; + if (blockTable.has_value()) + { + auto const& table = blockTable.value(); + TORCH_CHECK(table.is_cuda() && table.device() == device && table.scalar_type() == th::kInt32 && table.dim() == 2 + && table.stride(1) == 1 && table.size(0) >= numSeqs && table.size(1) > 0 && table.size(1) <= kInt32Max + && table.stride(0) >= table.size(1), + "block_table must be an int32 CUDA [sequences, max_blocks] tensor with unit column stride"); + TORCH_CHECK(tokensPerBlock > 0 && pageIndexScale > 0 && layerOffset >= 0, + "invalid block_table geometry: ", "tokens_per_block ", tokensPerBlock, ", page_index_scale ", + pageIndexScale, ", layer_offset ", layerOffset); + int64_t const strideFactor = pageIndexScale * tokensPerBlock; + int64_t const rowOffset = layerOffset * tokensPerBlock; + TORCH_CHECK(strideFactor <= kInt32Max && rowOffset <= kInt32Max, "block_table rows must be int32-addressable"); + params.remapBlockTable = table.data_ptr(); + params.remapBlockTableStride = table.stride(0); + params.remapMaxBlocks = static_cast(table.size(1)); + params.remapTokensPerBlock = static_cast(tokensPerBlock); + params.remapStrideFactor = static_cast(strideFactor); + params.remapLayerOffset = static_cast(rowOffset); + } + else + { + TORCH_CHECK(!gatherBeforeWait, "gather_before_wait needs block_table"); + } + + auto const stream = at::cuda::getCurrentCUDAStream(q.get_device()); + tk::invokeGlmSparseMlaFlashDecode(params, stream); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_sparse_mla_flash_decode(Tensor q, Tensor kv_pool, Tensor indices, Tensor seq_lens_kv, " + "Tensor mla_bmm1_scale, Tensor mla_bmm2_scale, Tensor v_b_proj, Tensor(a!) output, int seq_len_q, int top_k, " + "Tensor? block_table=None, int tokens_per_block=0, int page_index_scale=0, int layer_offset=0, " + "bool gather_before_wait=False) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_sparse_mla_flash_decode", &tensorrt_llm::torch_ext::glmSparseMlaFlashDecode); +} diff --git a/cpp/tensorrt_llm/thop/glmVerifyHeadAcceptOp.cpp b/cpp/tensorrt_llm/thop/glmVerifyHeadAcceptOp.cpp new file mode 100644 index 000000000000..cfeb6d2ea623 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmVerifyHeadAcceptOp.cpp @@ -0,0 +1,108 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +// Accepted-token counts of the strict acceptance of one-model speculative decoding, with the forced count of +// TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS applied (force_mode 0: none, 1: force_base, 2: force_base plus a pool +// draw below force_frac, capped at force_max, force_counter advanced by one). Writes the accepted-token rows into +// accepted_tokens when given. +at::Tensor glmVerifyHeadAccept(at::Tensor const& target_tokens, at::Tensor const& draft_tokens, int64_t num_contexts, + std::optional accepted_tokens, std::optional const& force_pool, + std::optional force_counter, int64_t force_mode, int64_t force_base, int64_t force_max, + double force_frac, int64_t counter_stride, int64_t slot_stride, int64_t pool_size) +{ + CHECK_TH_CUDA(target_tokens); + CHECK_CONTIGUOUS(target_tokens); + auto const targetType = target_tokens.scalar_type(); + TORCH_CHECK(targetType == at::ScalarType::Int || targetType == at::ScalarType::Long, + "glm_verify_head_accept expects int32 or int64 target tokens, got ", targetType); + CHECK_INPUT(draft_tokens, at::ScalarType::Int); + TORCH_CHECK(draft_tokens.dim() == 2, "glm_verify_head_accept expects draft tokens [num_gens, draft_len]"); + int64_t const numGens = draft_tokens.size(0); + int64_t const draftLen = draft_tokens.size(1); + int64_t const batchSize = num_contexts + numGens; + TORCH_CHECK(target_tokens.numel() >= num_contexts + numGens * (draftLen + 1), + "glm_verify_head_accept: target tokens hold one row per context and draft_len + 1 per generation request"); + CHECK_OPTIONAL_INPUT(accepted_tokens, at::ScalarType::Int); + TORCH_CHECK(!accepted_tokens.has_value() || accepted_tokens->numel() >= batchSize * (draftLen + 1), + "glm_verify_head_accept: accepted_tokens must be [batch, draft_len + 1]"); + TORCH_CHECK(force_mode >= 0 && force_mode <= 2, "glm_verify_head_accept: unknown force_mode ", force_mode); + bool const randomForce = force_mode == 2; + if (randomForce) + { + TORCH_CHECK(force_pool.has_value() && force_counter.has_value(), + "glm_verify_head_accept: force_mode 2 needs force_pool and force_counter"); + CHECK_INPUT(force_pool.value(), at::ScalarType::Float); + CHECK_INPUT(force_counter.value(), at::ScalarType::Long); + TORCH_CHECK(pool_size > 0 && (pool_size & (pool_size - 1)) == 0 && force_pool->numel() >= pool_size, + "glm_verify_head_accept: the pool size must be a power of two within force_pool"); + } + + at::Tensor numAccepted = at::empty({batchSize}, draft_tokens.options()); + + namespace tk = tensorrt_llm::kernels::glm_moe_dsa; + tk::VerifyHeadAcceptParams params; + params.targetTokens = target_tokens.data_ptr(); + params.targetTokensInt64 = targetType == at::ScalarType::Long; + params.draftTokens = draft_tokens.data_ptr(); + params.numAcceptedTokens = numAccepted.data_ptr(); + params.acceptedTokens = accepted_tokens.has_value() ? accepted_tokens->data_ptr() : nullptr; + params.numContexts = static_cast(num_contexts); + params.numGens = static_cast(numGens); + params.draftLen = static_cast(draftLen); + params.forceMode = static_cast(force_mode); + params.forceBase = static_cast(force_base); + params.forceMax = static_cast(force_max); + params.forceFrac = static_cast(force_frac); + params.forcePool = randomForce ? force_pool->data_ptr() : nullptr; + params.forceCounter = randomForce ? force_counter->data_ptr() : nullptr; + params.counterStride = counter_stride; + params.slotStride = slot_stride; + params.poolMask = pool_size - 1; + params.stream = at::cuda::getCurrentCUDAStream(draft_tokens.get_device()); + tk::invokeVerifyHeadAccept(params); + return numAccepted; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_verify_head_accept(Tensor target_tokens, Tensor draft_tokens, int num_contexts, " + "Tensor(a!)? accepted_tokens, Tensor? force_pool, Tensor(b!)? force_counter, int force_mode, int force_base, " + "int force_max, float force_frac, int counter_stride, int slot_stride, int pool_size) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_verify_head_accept", &tensorrt_llm::torch_ext::glmVerifyHeadAccept); +} diff --git a/cpp/tensorrt_llm/thop/moe/trtllmGen/fp4BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/moe/trtllmGen/fp4BlockScaleMoe.cpp index fac7cfb36e77..5333d5daa8dd 100644 --- a/cpp/tensorrt_llm/thop/moe/trtllmGen/fp4BlockScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/moe/trtllmGen/fp4BlockScaleMoe.cpp @@ -46,7 +46,11 @@ std::vector run_fp4_block_scale_moe_runner(torch::optional const routed_scaling_factor, int64_t const tile_tokens_dim, int64_t const routing_method_type, bool const do_finalize, btg::Dtype const dtype, MoeRunnerType& moe_runner, int64_t const moeConfigIndex, torch::optional const& topk_weights, torch::optional const& topk_ids, - torch::optional const& out_tensor = torch::nullopt) + torch::optional const& out_tensor = torch::nullopt, + std::optional const hidden_states_ready_event = std::nullopt, + torch::optional const& router_hidden_states = torch::nullopt, + torch::optional const& router_weight = torch::nullopt, + torch::optional const& router_workspace = torch::nullopt, bool const prefer_gemm2_k128 = false) { TORCH_CHECK(dtype == btg::Dtype::E4m3 || dtype == btg::Dtype::E2m1, "dtype can only be e4m3 or e2m1."); TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "Only SM100f is supported by FP4 block scale MOE"); @@ -154,6 +158,10 @@ std::vector run_fp4_block_scale_moe_runner(torch::optional(static_cast(hidden_states_ready_event.value())) + : nullptr; args.topk_weights = topk_weights.has_value() ? topk_weights.value().data_ptr() : nullptr; args.topk_ids = topk_ids.has_value() ? static_cast(topk_ids.value().data_ptr()) : nullptr; @@ -286,16 +294,50 @@ std::vector run_fp4_block_scale_moe_runner(torch::optional(), - expert_count_histogram.data_ptr(), total_num_padded_tokens.data_ptr(), - expanded_idx_to_permuted_idx.data_ptr(), nullptr, /*permuted_idx_to_expanded_idx.data_ptr(),*/ - permuted_idx_to_token_idx.data_ptr(), expert_weights_ptr, args.topk_ids, - num_tokens_per_expert.data_ptr(), cta_idx_xy_to_batch_idx.data_ptr(), - cta_idx_xy_to_mn_limit.data_ptr(), num_non_exiting_ctas.data_ptr(), args.mDtypeElt, - false /* use_routing_scales_on_input */, false /* use_deep_seek_fp8 */, - static_cast(routing_method_type), stream, dtypeRoutingLogits, args.mDtypeBias); + if (router_hidden_states.has_value()) + { + // The router GEMV runs in the routing kernel: routing_logits only gives the shape and is not read or written. + TORCH_CHECK(router_weight.has_value() && router_workspace.has_value(), + "router_hidden_states needs router_weight and router_workspace."); + TORCH_CHECK(!args.topk_ids + && static_cast(routing_method_type) == RoutingMethodType::DeepSeekV3 + && args.n_group <= 1 && routing_bias.has_value(), + "The fused router covers DeepSeekV3 routing without groups, with a routing bias, from the router inputs."); + auto const& hs = router_hidden_states.value(); + auto const& rw = router_weight.value(); + auto const& ws = router_workspace.value(); + TORCH_CHECK(hs.scalar_type() == at::ScalarType::BFloat16 && hs.dim() == 2 && hs.is_contiguous() + && hs.sizes()[0] == args.num_tokens && hs.sizes()[1] == args.hidden_size + && hs.device() == routing_device, + "router_hidden_states must be a contiguous bf16 [num_tokens, hidden_size] tensor on the routing device."); + TORCH_CHECK(rw.scalar_type() == at::ScalarType::BFloat16 && rw.dim() == 2 && rw.is_contiguous() + && rw.sizes()[0] == num_experts && rw.sizes()[1] == hs.sizes()[1] && rw.device() == routing_device, + "router_weight must be a contiguous bf16 [num_experts, hidden] tensor on the routing device."); + TORCH_CHECK(ws.scalar_type() == at::ScalarType::Int && ws.is_contiguous() + && ws.numel() == moe::dev::routing::routingCustom::FusedRouterWorkspaceWords + && ws.device() == routing_device, + "router_workspace must be a contiguous int32 tensor of ", + moe::dev::routing::routingCustom::FusedRouterWorkspaceWords, " words on the routing device."); + routing_runner.runFusedRouter(hs.data_ptr(), rw.data_ptr(), static_cast(hs.sizes()[1]), + ws.data_ptr(), args.routing_bias, args.num_tokens, args.num_experts, args.top_k, + args.local_expert_offset, args.local_num_experts, args.routed_scaling_factor, + total_num_padded_tokens.data_ptr(), expanded_idx_to_permuted_idx.data_ptr(), nullptr, + permuted_idx_to_token_idx.data_ptr(), expert_weights_ptr, cta_idx_xy_to_batch_idx.data_ptr(), + cta_idx_xy_to_mn_limit.data_ptr(), num_non_exiting_ctas.data_ptr(), stream, args.mDtypeBias); + } + else + { + routing_runner.run(args.routing_logits, args.routing_bias, args.num_tokens, args.num_experts, args.top_k, + /* num_fused_shared_expert */ 0, args.n_group, args.topk_group, args.local_expert_offset, + args.local_num_experts, args.routed_scaling_factor, expert_indexes.data_ptr(), + expert_count_histogram.data_ptr(), total_num_padded_tokens.data_ptr(), + expanded_idx_to_permuted_idx.data_ptr(), nullptr, /*permuted_idx_to_expanded_idx.data_ptr(),*/ + permuted_idx_to_token_idx.data_ptr(), expert_weights_ptr, args.topk_ids, + num_tokens_per_expert.data_ptr(), cta_idx_xy_to_batch_idx.data_ptr(), + cta_idx_xy_to_mn_limit.data_ptr(), num_non_exiting_ctas.data_ptr(), args.mDtypeElt, + false /* use_routing_scales_on_input */, false /* use_deep_seek_fp8 */, + static_cast(routing_method_type), stream, dtypeRoutingLogits, args.mDtypeBias); + } // // FC13 (gemm1) + FC2 (gemm2) @@ -486,6 +528,7 @@ std::vector run_fp4_block_scale_moe_runner(torch::optional(); args.output2_scales_scalar = output2_scales_scalar.data_ptr(); args.do_finalize = do_finalize; + args.prefer_gemm2_k128 = prefer_gemm2_k128; auto const workspace_sizes = moe_runner.getWorkspaceSizeInBytes(args, moeConfigIndex); @@ -589,7 +632,11 @@ class FP4BlockScaleMoeRunner : public torch::CustomClassHolder int64_t const local_expert_offset, int64_t const local_num_experts, std::optional const routed_scaling_factor, int64_t const routing_method_type, bool const do_finalize, std::vector moeConfigIndex, torch::optional const& topk_weights, - torch::optional const& topk_ids, torch::optional const& output = torch::nullopt) + torch::optional const& topk_ids, torch::optional const& output = torch::nullopt, + std::optional const hidden_states_ready_event = std::nullopt, + torch::optional const& router_hidden_states = torch::nullopt, + torch::optional const& router_weight = torch::nullopt, + torch::optional const& router_workspace = torch::nullopt, bool const prefer_gemm2_k128 = false) { int64_t tileN = moeConfigIndex[0]; int64_t config = moeConfigIndex[1]; @@ -639,7 +686,8 @@ class FP4BlockScaleMoeRunner : public torch::CustomClassHolder gemm2_weights_scale, gemm2_bias, output1_scales_scalar, output1_scales_gate_scalar, output2_scales_scalar, num_experts, top_k, n_group, topk_group, intermediate_size, local_expert_offset, local_num_experts, routed_scaling_factor, tileN, routing_method_type, do_finalize, mDtypeElt, *runner, configToRun, - topk_weights, topk_ids, output); + topk_weights, topk_ids, output, hidden_states_ready_event, router_hidden_states, router_weight, + router_workspace, prefer_gemm2_k128); } private: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index b0b6e78ef8d8..5233f575bd73 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -54,6 +54,7 @@ # Qwen3.8-Flash-Next publishes ``Qwen4ExpForConditionalGeneration``, and the # config loader's composite branch writes the same name back. "Qwen4ExpForConditionalGeneration": "models.qwen4_exp.routing", + "GlmMoeDsaForCausalLM": "models.glm_moe_dsa.routing", } diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/_op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/_op.py index 3dc75761b4a3..4baee1ba64b6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/_op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/_op.py @@ -88,6 +88,7 @@ class Arch(str, enum.Enum): """ SM_103 = "sm_103" + SM_107 = "sm_107" @dataclass(frozen=True) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py index b0c0e864d886..2e23c1c3c997 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py @@ -51,10 +51,10 @@ class Fp4Quantize(OpWrapper): # nvfp4_gemm in the dense and shared-expert MLPs, the linear one feeding the # routed-expert path, at its hidden size and its intermediate. qwen3.8 # (qwen4_exp) calls it only for its routed experts: linear scales at hidden - # 2560. A linear buffer has no row padding, so at a row count that is not a - # multiple of 128 it is shorter than the swizzled one; at a multiple of 128 - # the two are the same length and only the cells' byte order tells them - # apart. + # 2560; glm-5.2 likewise, at hidden 6144. A linear buffer has no row + # padding, so at a row count that is not a multiple of 128 it is shorter + # than the swizzled one; at a multiple of 128 the two are the same length + # and only the cells' byte order tells them apart. CELLS: tuple[Cell, ...] = ( Cell( why="decode, hidden 7168, swizzled scales -- the dense MLP's input", @@ -88,6 +88,10 @@ class Fp4Quantize(OpWrapper): why="qwen3.8 prefill at max_num_tokens 8192, hidden 2560, linear", spec=dict(shape=(8192, 2560), dtype=torch.bfloat16, vec=16, swizzled=False), ), + Cell( + why="glm-5.2 verify step: 8 rows at hidden 6144, linear -- its routed-expert input", + spec=dict(shape=(8, 6144), dtype=torch.bfloat16, vec=16, swizzled=False), + ), ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/__init__.py new file mode 100644 index 000000000000..43c6570013d1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GlmMoeDsaForCausalLM targets (GLM-5.2) and the routing between them.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/__init__.py new file mode 100644 index 000000000000..9bcf8ba8c02c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GLM-5.2-NVFP4 on 16 Rubin GPUs (sm107): tensor-parallel attention and experts over 16 ranks.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py new file mode 100644 index 000000000000..9f484bf95dee --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py @@ -0,0 +1,1973 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GLM-5.2-NVFP4 on 16 Rubin GPUs (sm107), tp16: model construction and forward. + +This target serves the generation side of a disaggregated deployment. Its +context server runs the built-in GLM-5.2 implementation (modeling_deepseekv3) +and hands over the MLA KV cache and the DSA indexer K cache it wrote, so this +target builds the same modules and loads the same weights: the caches it reads +are laid out exactly as they were written. Forwards that carry context rows +(engine warm-up, an aggregated deployment) run that same decoder stack in +``PrefillTarget``. Generation-only forwards -- every serving step of the +generation server -- run ``DecodeTarget``. With one-model MTP-Eagle, the draft +loop of a one-request generation step runs on ``DecodeTarget``'s kernels as +well (``MTPDraftWorker``). +""" + +from __future__ import annotations + +import contextlib +from dataclasses import dataclass +from typing import Any, List, Literal, Optional, Tuple + +import torch +from transformers import PretrainedConfig + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + AttentionInputType, + AttentionMetadata, +) +from tensorrt_llm._torch.attention.backends.sparse.dsa.module import forward_generation_sparse_attn +from tensorrt_llm._torch.attention.backends.sparse.dsa.params import DSABackendForwardArgs +from tensorrt_llm._torch.attention.backends.sparse.params import MTPIndexShareMetadata +from tensorrt_llm._torch.distributed import AllReduceFusionOp, AllReduceParams +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3ForCausalLM, DeepseekV3Model +from tensorrt_llm._torch.models.modeling_speculative import ( + SpecDecOneEngineForCausalLM, + _slice_spec_position_ids, +) +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.modules.swiglu import swiglu +from tensorrt_llm._torch.moe.fused_moe.trtllm_gen.kernel_inputs import ( + extract_routing_params, + get_data_or_none, + to_trtllm_gen_act_type, +) +from tensorrt_llm._torch.speculative import MTPEagleWorker, SpecMetadata +from tensorrt_llm._torch.speculative.eagle3 import _reset_mtp_index_share +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, +) +from tensorrt_llm._torch.utils import EventType + +from . import weights as _weights +from .ops.attention.mla_kv_half_append import mla_kv_half_append +from .ops.attention.sparse_mla_flash_decode import sparse_mla_flash_decode +from .ops.comm._mnnvl_workspace import ( + LamportExchangeWorkspace, + hop_supported, + hop_workspace, + publish_target, +) +from .ops.comm.mnnvl_allgather import mnnvl_allgather +from .ops.comm.mnnvl_hop import mnnvl_hop +from .ops.comm.mnnvl_moe_finalize_hop import mnnvl_moe_finalize_hop +from .ops.dsa._fp8_paged_mqa_logits import num_sms +from .ops.dsa._hist_topk import ( + MAX_SHARD, + TOPK, + get_hist_topk_workspace, + hist_topk_supported, + shard_capacity, +) +from .ops.dsa.convert_req_index_to_global import convert_req_index_to_global +from .ops.dsa.fp8_paged_mqa_logits import fp8_paged_mqa_logits +from .ops.dsa.hist_shard_meta import hist_shard_meta +from .ops.dsa.hist_topk import hist_topk +from .ops.dsa.indexer_k_tail import indexer_k_tail +from .ops.dsa.paged_mqa_logits_schedule import paged_mqa_logits_schedule +from .ops.dsa.stash_last_accepted_rows import stash_last_accepted_rows +from .ops.embedding.fused_embedding import fused_embedding +from .ops.gemm.dense_mlp_small_m import INTER_PER_RANK, dense_mlp_small_m +from .ops.gemm.draft_hidden_half import draft_hidden_half +from .ops.gemm.fused_a_gemm import INDEXER_OUT_FEATURES as FUSED_A_INDEXER_OUT_FEATURES +from .ops.gemm.fused_a_gemm import MAX_TOKENS as FUSED_A_MAX_TOKENS +from .ops.gemm.fused_a_gemm import OUT_FEATURES as FUSED_A_OUT_FEATURES +from .ops.gemm.fused_a_gemm import fused_a_gemm +from .ops.gemm.indexer_q_tail import indexer_q_tail +from .ops.gemm.low_m_gemm import low_m_gemm +from .ops.gemm.o_proj_publish import o_proj_publish +from .ops.gemm.q_b_fp8_rope import q_b_fp8_rope +from .ops.gemm.router_gemm import router_gemm +from .ops.moe.lowm_bf16_moe import glm_lowm_bf16_moe +from .ops.moe.routed_experts import ( + FUSED_ROUTER_MAX_TOKENS, + make_fused_router_workspace, + routed_experts, +) +from .ops.spec._draft_exchange import HIDDEN_SIZE as DRAFT_HIDDEN_SIZE +from .ops.spec._draft_exchange import DraftExchangeWorkspace, ehproj_table_supports_tp +from .ops.spec._draft_token_table import build_draft_token_table +from .ops.spec.draft_ehproj_table import glm_draft_ehproj_table +from .ops.spec.draft_glue_rmsnorm_cat import glm_draft_glue_rmsnorm_cat +from .ops.spec.draft_sample_glue import glm_draft_sample_glue +from .ops.spec.gather_last_accepted_rows import glm_gather_last_accepted_rows +from .ops.spec.verify_head_accept import glm_verify_head_accept + +# The deployment this target is written for; route() selects it on the same +# facts, and the constructor states them again for a direct construction. +_COMPUTE_CAPABILITY = (10, 7) +_TP = 16 +# One Lamport buffer of the verify logits' all-gather: every rank's BF16 vocabulary shard of +# up to 8 rows, with room to spare. +_LOGITS_EXCHANGE_BYTES = 16 << 20 + + +class PrefillTarget: + """Forwards that carry context rows: the decoder stack the context server runs.""" + + def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: + self.core = core + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, + **kwargs, + ) -> torch.Tensor: + return DeepseekV3Model.forward( + self.core, + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + **kwargs, + ) + + +@dataclass(frozen=True, eq=False) +class _DraftState: + """What the MTP draft steps hold for the engine's life (``DecodeTarget.build_draft``).""" + + # The MTP layer (the built-in module at layer num_hidden_layers) and the LM head it shares. + layer: Any + lm_head: Any + # Detached views of the embedding table and the enorm / hnorm weights for the CuTe kernels. + embed_table: torch.Tensor + enorm_weight: torch.Tensor + hnorm_weight: torch.Tensor + # This rank's columns [lo, hi) of eh_proj's concatenated input (eh_proj is row-parallel), + # and whether draft step 0's glue kernel launches with PDL. + window: Tuple[int, int] + pdl: bool + # What fused-A prefetches into L2 on the draft steps that run no indexer. + l2_after_fused_a: List[torch.Tensor] + # The residual of draft step 0's eh_proj hop: zeros, so its residual output is the sum. + zero_residual: torch.Tensor + # glm_draft_ehproj_table's operands: the token half of eh_proj for the tokens of this rank's + # vocabulary shard, and this rank's rows of eh_proj's hidden half. + token_table: torch.Tensor + hidden_rows: torch.Tensor + # The low-M BF16 MoE's experts and routing constants. + w3_w1: torch.Tensor + w2: torch.Tensor + top_k: int + routed_scaling_factor: float + # The post-MoE hop's routing: one slot of weight 1 per row, holding that row (-1 elsewhere). + expert_weights: torch.Tensor + expert_rows: torch.Tensor + # Draft step 0's top-k of the request's last accepted row, and the request's index. + topk_stash: torch.Tensor + request_rows: torch.Tensor + # The exchange workspaces of glm_draft_ehproj_table and glm_draft_sample_glue. + table_exchange: DraftExchangeWorkspace + sample_exchange: DraftExchangeWorkspace + # The side stream that stashes that top-k and maps it to KV pool rows, its fork and ready + # events, and the events of the hidden half, which runs on the MTP layer's MoE aux stream. + stream: torch.cuda.Stream + selection_fork: torch.cuda.Event + selection_ready: torch.cuda.Event + half_start: torch.cuda.Event + half_done: torch.cuda.Event + + +class DecodeTarget: + """Generation-only forwards: the serving step of the generation server. + + Routed to only when the batch has no context rows, which the core decides + per step; nothing in this class reads the phase back. + + A step of up to 8 rows (one request and its 7 draft tokens at c1) runs on + the target's kernels, layer by layer: + + - The embedding is one lookup in the whole BF16 table every rank holds. + - Attention: q_a, the compressed KV and its rope part come from one GEMM over + the stored kv_a_proj_with_mqa weight. q_b, with k_b absorbed and the q_a + norm folded, writes the RoPE'd FP8 query in its epilogue, while on the aux + stream one kernel normalizes, ropes and appends the KV half to the FP8 + cache. The flash decode attends over the 2048 selected rows at 4 heads with + v_b folded in, and o_proj's epilogue publishes this rank's shard into the + all-reduce workspace. + - The DSA indexer (on the layers that run one) shares fused-A: its K and + weight rows are fused-A's last 160 columns. Its Q comes from wq_b (q_a norm + folded) with an epilogue that ropes, quantizes and scales; its K tail + normalizes, ropes, quantizes and writes the indexer K cache on a side + stream. Each rank scores its 1/16 shard of the context's keys, and the ranks + agree on the exact top-2048 by exchanging histograms. The layers without an + indexer reuse the last selection, which the flash decode maps to cache rows + itself and starts gathering before its wait. + - Every all-reduce is a one-shot MNNVL hop that adds the residual and applies + the following RMSNorm: post_attention_layernorm after attention, and after + the MLP or MoE the next layer's input_layernorm (the model's final norm + after the last layer). + - The dense layers 0-2 run this rank's 768 intermediate columns as two + weight-streaming kernels. + - The MoE layers run the shared expert as two split-K GEMMs around SwiGLU on + the aux stream, and the routed experts as one trtllm-gen call that runs the + router GEMV in its routing kernel, joins the NVFP4 input (quantized on its + own stream) right before FC1, runs FC2 on the K = 128 kernel and leaves the + finalize step to the post-MoE hop. + + The verify logits come from the LM head on this rank's vocabulary shard + (split-K GEMM) and one MNNVL all-gather that widens them to FP32. + + A larger batch runs the built-in modules in the same order. + + The MTP draft steps of a one-request step run here too (``draft_first_step``, + ``draft_next_step``, ``draft_head``, ``draft_ehproj_table``, ``draft_sample_glue``), + on state that ``build_draft`` makes once: + + - Draft step 0 runs on the verify rows. One kernel writes this rank's window of + eh_proj's input (embedding gather, enorm, hnorm, concat, column split), eh_proj + runs on it as a split-K GEMM, and its hop applies input_layernorm with a zero + residual. The MTP layer's attention is a verify full-indexer layer's; each + request's top-k of its last accepted row is kept, as KV pool rows of the MTP + layer, for the later steps (on a side stream). After the post-attention hop each + request keeps the row of its last accepted token. + - The later steps take eh_proj's output and its input_layernorm from the previous + step's table kernel. They run no indexer: they attend to draft step 0's + selection, or, with a context short enough that the built-in metadata skips the + indexer, to every position up to their own. + - Every step's MoE runs on that one row: the router GEMV and the low-M BF16 experts, + the shared expert on the aux stream, and the hop that adds both and the residual + and applies shared_head.norm. + - The draft LM head runs on this rank's vocabulary shard. Before every step but the + last, the next step's eh_proj hidden half starts first on the aux stream, and + ``glm_draft_ehproj_table`` exchanges the ranks' argmax and writes the draft + token, the next step's eh_proj output from a token table built at load, its + positions and the sequence and KV lengths. The last step's + ``glm_draft_sample_glue`` writes its token and the next-new-tokens rows. + """ + + def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: + self.core = core + config = core.model_config.pretrained_config + self.mapping = core.model_config.mapping + self.hidden_size = config.hidden_size + self.index_topk = config.index_topk + layers = core.layers[: core.num_hidden_layers] + for layer in layers: + assert layer.allreduce is not None and layer.next_layer_layernorm is not None + self._check_attention(layer) + if isinstance(layer.mlp, GatedMLP): + assert ( + layer.mlp_tp_size == _TP + and not layer.disable_attn_allreduce + and not layer.fusion_config.PRE_MLP_FUSION + and layer.fusion_config.POST_MLP_FUSION + and layer.post_attention_layernorm.nvfp4_scale is None + ), f"layer {layer.layer_idx}: unexpected dense-layer configuration" + else: + assert ( + layer.disable_attn_allreduce + and layer.fusion_config.PRE_MOE_FUSION + and layer.fusion_config.POST_MOE_FUSION + and layer.mapping.is_multi_node() + ), f"layer {layer.layer_idx}: unexpected MoE-layer configuration" + self._check_shared_expert(layer) + backend = getattr(layer.mlp.experts, "backend", layer.mlp.experts) + assert ( + backend.has_nvfp4 + and backend.scaling_vector_size == 16 + and getattr(backend, "fc31_act_scale", None) is None + and not backend.bias + and backend.w3_w1_weight.shape[-1] * 2 == config.hidden_size + ), f"layer {layer.layer_idx}: routed experts are not unpadded TRTLLM-Gen NVFP4" + + # The attention operands the kernels read, per layer index (_fold_attention_weights). + self._q_b_absorbed = {} + self._fused_a_indexer_weight = {} + self._wq_b_folded = {} + self._l2_after_fused_a = {} + for layer in layers: + self._fold_attention_weights(layer) + # The key-sharded top-k exchanges histograms over the TP group's MNNVL workspace + # (created here, on every rank in the same order). + self._hist_workspace = ( + get_hist_topk_workspace(self.mapping) if hist_topk_supported(self.mapping) else None + ) + self._hist_buffers = None + self._logits_workspace = LamportExchangeWorkspace(self.mapping, _LOGITS_EXCHANGE_BYTES) + self._k_tail_stream = torch.cuda.Stream() + self._k_tail_fork = torch.cuda.Event() + self._k_tail_done = torch.cuda.Event() + + device = core.embed_tokens.weight.device + # The attention scales the KV-half kernel writes and the flash decode reads. + self._bmm1_scale = torch.empty(2, dtype=torch.float32, device=device) + self._bmm2_scale = torch.empty(1, dtype=torch.float32, device=device) + self._unit_scale = torch.ones(1, dtype=torch.float32, device=device) + + moe_layers = [layer for layer in layers if not isinstance(layer.mlp, GatedMLP)] + # One exchange workspace per router call site: its launches must run in stream order. + self._router_workspace = { + layer.layer_idx: make_fused_router_workspace(device) for layer in moe_layers + } + # The routing kernel runs the router GEMV, so the logits it is handed only give the + # shape; autotuner profiling routes from them, so they hold random values. + self._router_logits = torch.randn( + FUSED_ROUTER_MAX_TOKENS, config.n_routed_experts, device=device + ) + self._quant_stream = torch.cuda.Stream() + self._quant_fork = torch.cuda.Event() + self._quant_ready = torch.cuda.Event() + self._fast_rows = {} + # Whether the last verify forward put the indexer layers on this rank's key shard; draft + # step 0 of the same step runs its indexer on that shard. + self._keys_sharded = False + self._draft: Optional[_DraftState] = None + + def _check_attention(self, layer) -> None: + """Assert the layer's MLA weights and KV cache are the layout the kernels take.""" + mla = layer.self_attn + assert not mla.apply_rotary_emb and not mla.mapping.has_cp_helix() + assert ( + tuple(mla.kv_a_proj_with_mqa.weight.shape) + == ( + FUSED_A_OUT_FEATURES, + self.hidden_size, + ) + and mla.kv_a_proj_with_mqa.weight.dtype == torch.bfloat16 + ), f"layer {layer.layer_idx}: kv_a_proj_with_mqa is not the BF16 [q_a; kv_a] weight" + assert ( + getattr(mla.mqa, "has_fp8_kv_cache", False) + and mla.kv_cache_dtype != "fp8_ds_mla" + and mla.k_b_proj_trans.dtype == torch.bfloat16 + and mla.v_b_proj.dtype == torch.bfloat16 + and mla.v_b_proj.is_contiguous() + and mla.q_b_proj.weight.dtype == torch.bfloat16 + and mla.o_proj.weight.dtype == torch.bfloat16 + ), f"layer {layer.layer_idx}: unexpected MLA weights or KV cache layout" + assert layer.input_layernorm.nvfp4_scale is None, "kv_a_proj is not NVFP4" + + @staticmethod + def _check_shared_expert(layer) -> None: + """Assert the MoE layer's shared expert is the plain BF16 SwiGLU MLP the kernels run.""" + shared = layer.mlp.shared_experts + assert ( + shared is not None + and layer.mlp.shared_output_scale is None + and shared._is_plain_swiglu() + and shared.swiglu_limit is None + and shared.gate_up_proj.weight.dtype == torch.bfloat16 + and shared.down_proj.weight.dtype == torch.bfloat16 + ), f"layer {layer.layer_idx}: the shared expert is not a plain BF16 SwiGLU MLP" + + def _fold_attention_weights(self, layer) -> None: + """The layer's attention operands in the forms the kernels read, by layer index.""" + mla = layer.self_attn + index = layer.layer_idx + # q_b with k_b absorbed and the q_a norm weight folded: for head h, rows + # [h * 576, h * 576 + 512) are k_b_trans[h] @ W_q_b[h, :qk_nope] and the next 64 rows + # W_q_b[h, qk_nope:], every column times gamma_q_a; FP64, rounded to BF16 once. + qk_head_dim = mla.qk_nope_head_dim + mla.qk_rope_head_dim + w = mla.q_b_proj.weight.view(mla.num_heads_tp, qk_head_dim, mla.q_lora_rank).double() + nope = torch.bmm(mla.k_b_proj_trans.double(), w[:, : mla.qk_nope_head_dim]) + absorbed = torch.cat([nope, w[:, mla.qk_nope_head_dim :]], 1) + absorbed = absorbed * mla.q_a_layernorm.weight.double() + self._q_b_absorbed[index] = ( + absorbed.to(torch.bfloat16).reshape(-1, mla.q_lora_rank).contiguous() + ) + + # A layer that runs the indexer gets its fused-A weight: kv_a_proj_with_mqa's rows, + # then the indexer's wk and weights_proj rows (BF16), in one tensor whose first rows + # the built-in kv_a_proj and the indexer's projections then view, so nothing is stored + # twice. Its indexer Q comes from wq_b with the q_a norm weight folded into its columns. + indexer = mla.mqa.indexer + if indexer is not None: + assert ( + indexer.wk.weight.dtype == torch.bfloat16 + and indexer.weights_proj.weight.dtype == torch.bfloat16 + ), "the indexer's wk and weights_proj must be BF16 (TRTLLM_DSA_INDEXER_BF16=1)" + kv_a = mla.kv_a_proj_with_mqa.weight + wk = indexer.wk.weight + extended = torch.cat([kv_a, wk, indexer.weights_proj.weight]).contiguous() + assert extended.shape[0] == FUSED_A_INDEXER_OUT_FEATURES + mla.kv_a_proj_with_mqa.weight.data = extended[: kv_a.shape[0]] + indexer.wk.weight.data = extended[kv_a.shape[0] : kv_a.shape[0] + wk.shape[0]] + indexer.weights_proj.weight.data = extended[kv_a.shape[0] + wk.shape[0] :] + self._fused_a_indexer_weight[index] = extended + self._wq_b_folded[index] = ( + indexer.wq_b.weight.double() * mla.q_a_layernorm.weight.double() + ).to(torch.bfloat16) + + # What fused-A prefetches into L2 while it waits for its input: the weights of the + # attention GEMMs that follow it. + prefetch = [self._q_b_absorbed[index], mla.v_b_proj, mla.o_proj.weight] + if index in self._wq_b_folded: + prefetch.insert(1, self._wq_b_folded[index]) + self._l2_after_fused_a[index] = [t for t in prefetch if t.is_contiguous()] + + @torch.no_grad() + def build_draft(self, layer, lm_head, weights: _weights.DraftWeights) -> None: + """Build what the MTP draft steps hold for the engine's life, once. + + ``layer`` is the MTP layer and ``lm_head`` the vocabulary-parallel LM head it shares; + ``weights`` holds the checkpoint tensors the token table and the low-M MoE take. + Collective: every rank creates the two draft exchange workspaces here, in the same + order. The draft steps after the first reuse draft step 0's selection, as the built-in + indexer does under MTP index share; without that share, or where the kernels do not + serve one row, nothing is built and the built-in draft loop runs. + """ + if self._draft is not None: + return + core = self.core + config = core.model_config.pretrained_config + mapping = self.mapping + tp_size, tp_rank = mapping.tp_size, mapping.tp_rank + hidden = self.hidden_size + mla = layer.self_attn + indexer = mla.mqa.indexer + if ( + indexer is None + or not indexer.mtp_index_share + or not ehproj_table_supports_tp(tp_size) + or not hop_supported(1, hidden, tp_size) + ): + return + self._check_attention(layer) + self._check_shared_expert(layer) + columns = 2 * hidden // tp_size + norms = (layer.enorm, layer.hnorm, layer.input_layernorm, layer.shared_head.norm) + assert ( + hidden == DRAFT_HIDDEN_SIZE + and layer.eh_proj.weight.dtype == torch.bfloat16 + and tuple(layer.eh_proj.weight.shape) == (hidden, columns) + and weights.eh_proj.dtype == torch.bfloat16 + and tuple(weights.eh_proj.shape) == (hidden, 2 * hidden) + and all( + norm.weight.dtype == torch.bfloat16 and tuple(norm.weight.shape) == (hidden,) + for norm in norms + ) + and layer.enorm.variance_epsilon == layer.hnorm.variance_epsilon + ), "the MTP layer's eh_proj and norms are not the BF16 GLM-5.2 layout" + experts = config.n_routed_experts + inter = config.moe_intermediate_size // tp_size + gate = layer.mlp.gate + assert ( + config.n_group == 1 + and gate.weight.dtype == torch.bfloat16 + and tuple(gate.weight.shape) == (experts, hidden) + and tuple(weights.w3_w1.shape) == (experts, 2 * inter, hidden) + and tuple(weights.w2.shape) == (experts, hidden, inter) + ), "the MTP layer's MoE is not the one-group BF16 MoE of the low-M kernels" + embed_table = core.embed_tokens.weight + vocab_shard = lm_head.weight.shape[0] + assert ( + lm_head.weight.dtype == torch.bfloat16 + and vocab_shard * tp_size == lm_head.num_embeddings == embed_table.shape[0] + ), "the draft LM head is not an unpadded BF16 vocabulary shard" + + # The collective part first, before anything that can fail on one rank only. + table_exchange = DraftExchangeWorkspace.for_ehproj_table(mapping) + sample_exchange = DraftExchangeWorkspace.for_sample_glue(mapping) + + self._fold_attention_weights(layer) + device = embed_table.device + enorm_weight = layer.enorm.weight.detach() + token_table, hidden_rows = build_draft_token_table( + weights.eh_proj, + embed_table.detach(), + enorm_weight, + layer.enorm.variance_epsilon, + tp_rank, + tp_size, + vocab_shard, + ) + top_k = config.num_experts_per_tok + expert_weights = torch.zeros((1, top_k), dtype=torch.bfloat16) + expert_weights[0, 0] = 1 + expert_rows = torch.full((1, top_k), -1, dtype=torch.int32) + expert_rows[0, 0] = 0 + l2_after_fused_a = [self._q_b_absorbed[layer.layer_idx], mla.v_b_proj, mla.o_proj.weight] + self._draft = _DraftState( + layer=layer, + lm_head=lm_head, + embed_table=embed_table.detach(), + enorm_weight=enorm_weight, + hnorm_weight=layer.hnorm.weight.detach(), + window=(tp_rank * columns, (tp_rank + 1) * columns), + pdl=get_env_enable_pdl(), + l2_after_fused_a=[t for t in l2_after_fused_a if t.is_contiguous()], + zero_residual=torch.zeros( + (FUSED_A_MAX_TOKENS, hidden), dtype=torch.bfloat16, device=device + ), + token_table=token_table, + hidden_rows=hidden_rows, + w3_w1=weights.w3_w1, + w2=weights.w2, + top_k=top_k, + routed_scaling_factor=config.routed_scaling_factor, + expert_weights=expert_weights.to(device), + expert_rows=expert_rows.to(device), + topk_stash=torch.empty((1, self.index_topk), dtype=torch.int32, device=device), + request_rows=torch.zeros(1, dtype=torch.int32, device=device), + table_exchange=table_exchange, + sample_exchange=sample_exchange, + stream=torch.cuda.Stream(), + selection_fork=torch.cuda.Event(), + selection_ready=torch.cuda.Event(), + half_start=torch.cuda.Event(), + half_done=torch.cuda.Event(), + ) + + def _fast(self, rows: int) -> bool: + """Whether every kernel of the step serves this many rows (decided once per shape).""" + fast = self._fast_rows.get(rows) + if fast is None: + fast = ( + rows <= FUSED_A_MAX_TOKENS + and rows <= FUSED_ROUTER_MAX_TOKENS + and bool( + torch.ops.trtllm.glm_dense_mlp_small_m_supported( + rows, self.hidden_size, INTER_PER_RANK + ) + ) + and hop_supported(rows, self.hidden_size, self.mapping.tp_size) + ) + self._fast_rows[rows] = fast + return fast + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, + **kwargs, + ) -> torch.Tensor: + core = self.core + layers = core.layers[: core.num_hidden_layers] + if inputs_embeds is None: + inputs_embeds = fused_embedding(input_ids, core.embed_tokens.weight) + residual = inputs_embeds + hidden_states = layers[0].input_layernorm(inputs_embeds) + if not self._fast(hidden_states.shape[0]): + self._keys_sharded = False + for layer in layers: + hidden_states, residual = self._builtin_layer( + layer, position_ids, hidden_states, residual, attn_metadata + ) + return hidden_states + sharded = self._shard_indexer_keys(attn_metadata) + self._keys_sharded = sharded + for layer in layers: + attn_out = self._attention(layer, position_ids, hidden_states, attn_metadata, sharded) + hidden_states, residual = self._post_attention_hop(layer, attn_out, residual) + if isinstance(layer.mlp, GatedMLP): + mlp = layer.mlp + # This rank's partial sum; the hop adds the 16 ranks'. + partial = dense_mlp_small_m( + hidden_states, mlp.gate_up_proj.weight, mlp.down_proj.weight + ) + hidden_states, residual = self._hop(layer.next_layer_layernorm, partial, residual) + else: + hidden_states, residual = self._moe(layer, hidden_states, residual) + return hidden_states + + def logits(self, lm_head, hidden_states: torch.Tensor) -> Optional[torch.Tensor]: + """FP32 logits over the whole vocabulary of the rows, or None for a batch the kernels + do not serve: this rank's vocabulary shard from the LM head, then the all-gather.""" + rows = hidden_states.shape[0] + weight = lm_head.weight + if not self._fast(rows) or weight.shape[0] * self.mapping.tp_size != lm_head.num_embeddings: + return None + shard = low_m_gemm( + hidden_states.contiguous(), weight, hidden_states.new_empty(rows, weight.shape[0]) + ) + return mnnvl_allgather( + shard, + self._logits_workspace.comm_buffer, + self._logits_workspace.buffer_flags, + self.mapping.tp_size, + ) + + def drafts(self, rows: int) -> bool: + """Whether this step's MTP draft steps run here: the draft state exists, this step's + verify forward put the indexer on this rank's key shard (one request; draft step 0 + reuses it) and the kernels serve draft step 0's ``rows``.""" + return self._draft is not None and self._keys_sharded and self._fast(rows) + + def draft_first_step( + self, + attn_metadata, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + batch_indices: torch.Tensor, + num_accepted: torch.Tensor, + draft_len: int, + ): + """MTP draft step 0 over the verify rows, through the MTP layer. + + ``input_ids`` are the tokens the verify step accepted at every row, ``hidden_states`` + the target's output rows, ``num_accepted`` the accepted count of each request + (``batch_indices``). Returns ``(hidden, gather_ids, pool_rows)``: the MTP layer's + output (shared_head.norm applied) at each request's last accepted row; that row's index + among the verify rows; and the selection the later steps attend to, as KV pool rows, + or None where they attend to every position (a short context) or no step follows. + """ + draft = self._draft + layer = draft.layer + rows = hidden_states.shape[0] + col_lo, col_hi = draft.window + eh_proj_input = glm_draft_glue_rmsnorm_cat( + draft.embed_table, + input_ids, + hidden_states, + draft.enorm_weight, + draft.hnorm_weight, + layer.enorm.variance_epsilon, + col_lo, + col_hi, + draft.pdl, + ) + partial = low_m_gemm( + eh_proj_input, layer.eh_proj.weight, eh_proj_input.new_empty((rows, self.hidden_size)) + ) + # A zero residual: the residual output is eh_proj's output, the sum of the partials. + hidden_states, residual = self._hop( + layer.input_layernorm, partial, draft.zero_residual[:rows] + ) + attn_out = self._attention(layer, position_ids, hidden_states, attn_metadata, True) + pool_rows = None + if draft_len > 1 and not attn_metadata.skip_indexer_for_gen_reqs: + pool_rows = self._stash_selection(attn_metadata, num_accepted, draft_len + 1) + hidden_states, residual = self._post_attention_hop(layer, attn_out, residual) + gather_ids, hidden_states, residual = glm_gather_last_accepted_rows( + hidden_states, batch_indices, num_accepted, draft_len + 1, 0, second=residual + ) + return self._draft_moe(hidden_states, residual), gather_ids, pool_rows + + def draft_next_step( + self, + attn_metadata, + eh_proj_output: Tuple[torch.Tensor, torch.Tensor], + pool_rows: Optional[torch.Tensor], + join_selection: bool, + ) -> torch.Tensor: + """An MTP draft step after the first, one row per request, through the MTP layer. + + ``eh_proj_output`` is the previous step's ``(input_layernorm(x), x)`` for this step's + eh_proj output x. The layer runs no indexer. Its rows attend to ``pool_rows``, draft + step 0's selection (``join_selection``, on the first such step, joins the side stream + that made it); with a context short enough that the built-in metadata skips the + indexer, to every position up to their own, the dense selection that metadata lists + (refreshed for this step by the loop). Returns the MTP layer's output rows. + """ + draft = self._draft + layer = draft.layer + hidden_states, residual = eh_proj_output + if attn_metadata.skip_indexer_for_gen_reqs: + positions = attn_metadata.topk_indices_buffer[: hidden_states.shape[0]] + attn_out = self._attention( + layer, None, hidden_states, attn_metadata, True, positions=positions + ) + else: + if join_selection: + torch.cuda.current_stream().wait_event(draft.selection_ready) + attn_out = self._attention( + layer, None, hidden_states, attn_metadata, True, pool_rows=pool_rows + ) + hidden_states, residual = self._post_attention_hop(layer, attn_out, residual) + return self._draft_moe(hidden_states, residual) + + def draft_head(self, hidden_states: torch.Tensor, next_step: bool): + """The draft LM head on this rank's vocabulary shard: ``(logits, hidden_half)``. + + With ``next_step`` the next step's eh_proj hidden half is enqueued first, on the aux + stream: this rank's FP32 share of ``W_eh[:, H:] @ hnorm(hidden_states)``, hnorm inside + the split-K GEMM (``draft_ehproj_table`` joins it); else ``hidden_half`` is None. + """ + draft = self._draft + layer = draft.layer + rows = hidden_states.shape[0] + hidden_half = None + if next_step: + hidden_half = torch.empty( + (rows, draft.hidden_rows.shape[0]), + dtype=torch.float32, + device=hidden_states.device, + ) + aux_stream = layer.mlp.aux_stream + draft.half_start.record() + with torch.cuda.stream(aux_stream): + aux_stream.wait_event(draft.half_start) + draft_hidden_half( + hidden_states, + draft.hidden_rows, + hidden_half, + draft.hnorm_weight, + layer.hnorm.variance_epsilon, + ) + draft.half_done.record() + weight = draft.lm_head.weight + logits = low_m_gemm(hidden_states, weight, hidden_states.new_empty(rows, weight.shape[0])) + return logits, hidden_half + + def draft_ehproj_table( + self, + logits: torch.Tensor, + hidden_half: torch.Tensor, + draft_tokens: torch.Tensor, + step: int, + positions: torch.Tensor, + gather_ids: Optional[torch.Tensor], + seq_lens: torch.Tensor, + kv_lens: torch.Tensor, + rewind_accepted: Optional[torch.Tensor], + ): + """The end of draft step ``step`` when another follows, in one kernel per rank. + + The ranks exchange their argmax with its token-table row and their share of the hidden + half (joined here). Every rank writes the draft token into ``draft_tokens[:, step]``, + the next step's eh_proj output x and ``input_layernorm(x)``, and its positions: + ``positions`` (at ``gather_ids`` on draft step 0) plus one. On draft step 0 + ``seq_lens`` becomes 1 and ``kv_lens`` drops the rejected drafts + (``rewind_accepted``); afterwards ``kv_lens`` grows by one. Returns + ``((normed, x), positions)`` for the next step. + """ + draft = self._draft + norm = draft.layer.input_layernorm + rows = logits.shape[0] + normed = logits.new_empty((rows, self.hidden_size)) + residual = logits.new_empty((rows, self.hidden_size)) + next_positions = positions.new_empty((rows,)) + first = step == 0 + torch.cuda.current_stream().wait_event(draft.half_done) + glm_draft_ehproj_table( + logits, + draft.table_exchange.buffer, + draft.table_exchange.flags, + self.mapping.tp_size, + draft_tokens, + step, + hidden_half, + draft.token_table, + norm.weight, + norm.variance_epsilon, + normed, + residual, + positions, + next_positions, + gather_ids, + seq_lens, + first, + kv_lens, + 2 if first else 1, + rewind_accepted, + draft_tokens.shape[1], + None, + [], + None, + None, + None, + None, + None, + ) + return (normed, residual), next_positions + + def draft_sample_glue( + self, + logits: torch.Tensor, + draft_tokens: torch.Tensor, + step: int, + seq_lens: torch.Tensor, + kv_lens: torch.Tensor, + rewind_accepted: Optional[torch.Tensor], + accepted_tokens: torch.Tensor, + batch_indices: torch.Tensor, + num_accepted: torch.Tensor, + ) -> torch.Tensor: + """The end of the last draft step ``step``, in one kernel per rank. + + The ranks exchange their argmax; every rank writes the draft token into + ``draft_tokens[:, step]``, the lengths as ``draft_ehproj_table`` does, and returns the + next-new-tokens rows: each request's last accepted token, then its draft tokens. + """ + draft = self._draft + rows = logits.shape[0] + draft_len = draft_tokens.shape[1] + next_new_tokens = draft_tokens.new_empty((rows, draft_len + 1)) + first = step == 0 + glm_draft_sample_glue( + logits, + draft.sample_exchange.buffer, + draft.sample_exchange.flags, + self.mapping.tp_size, + draft_tokens, + step, + # No eh_proj input window and no positions: no step follows. + None, + None, + None, + None, + 0.0, + 0, + 0, + None, + None, + None, + None, + seq_lens, + first, + kv_lens, + 2 if first else 1, + rewind_accepted, + draft_len, + None, + [], + None, + None, + None, + None, + None, + accepted_tokens, + batch_indices, + num_accepted, + next_new_tokens, + ) + return next_new_tokens + + def _stash_selection(self, attn_metadata, num_accepted: torch.Tensor, next_n: int): + """The selection the MTP draft steps after the first attend to, made on the side + stream after draft step 0's attention: each request's top-k of its last accepted row + (the MTP index share), as KV pool rows of the MTP layer. The first later step joins + it (``selection_ready``).""" + draft = self._draft + manager = attn_metadata.kv_cache_manager + requests = num_accepted.shape[0] + tokens_per_block = manager.tokens_per_block + local_layer = draft.layer.self_attn.mqa.get_local_layer_idx(attn_metadata) + page_index_scale, layer_offset = manager.get_primary_pool_page_index_params(local_layer) + main_stream = torch.cuda.current_stream() + draft.selection_fork.record() + with torch.cuda.stream(draft.stream): + draft.stream.wait_event(draft.selection_fork) + stashed = stash_last_accepted_rows( + self._hist_buffers["topk"][: requests * next_n], + num_accepted, + draft.topk_stash[:requests], + next_n, + ) + pool_rows = convert_req_index_to_global( + draft.request_rows[:requests], + attn_metadata.block_table[:requests], + stashed, + tokens_per_block, + self.index_topk, + page_index_scale * tokens_per_block, + layer_offset, + ) + draft.selection_ready.record() + # Allocated on the side stream and read on this one: outside graph capture the memory + # must outlive this stream's reads (a captured graph keeps it). + if not torch.cuda.is_current_stream_capturing(): + pool_rows.record_stream(main_stream) + return pool_rows + + def _draft_moe(self, hidden_states: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: + """The MTP layer's BF16 MoE on one row and the post-MoE hop, which adds the routed and + shared outputs (the routed row as the one slot of weight 1) and the residual and + applies shared_head.norm: the normalized row.""" + draft = self._draft + layer = draft.layer + mlp = layer.mlp + gate = mlp.gate + + def routed() -> torch.Tensor: + router_logits = router_gemm(hidden_states, gate.weight) + return glm_lowm_bf16_moe( + hidden_states, + router_logits, + gate.e_score_correction_bias, + draft.w3_w1, + draft.w2, + draft.top_k, + draft.routed_scaling_factor, + ) + + routed_output, shared_output = maybe_execute_in_parallel( + routed, + lambda: self._shared_expert(mlp, hidden_states), + mlp.event_dict[EventType.Main], + mlp.event_dict[EventType.MoeShared], + mlp.aux_stream, + ) + norm = layer.shared_head.norm + comm_buffer, buffer_flags = hop_workspace( + self.mapping, hidden_states.shape[0], self.hidden_size + ) + normed, _ = mnnvl_moe_finalize_hop( + routed_output, + draft.expert_weights, + draft.expert_rows, + shared_output, + residual, + norm.weight, + norm.variance_epsilon, + comm_buffer, + buffer_flags, + ) + return normed + + def _shard_indexer_keys(self, attn_metadata) -> bool: + """Point this rank's indexer at its 1/16 shard of the request's keys for this step: + the shard's length and block table, and the logits kernel's schedule over it. + + Applies to one generation request whose shard fits the top-k kernel; returns whether it + did, and otherwise every indexer layer scores the whole context with the built-in + indexer.""" + manager = attn_metadata.kv_cache_manager + if self._hist_workspace is None or attn_metadata.num_seqs != 1: + return False + if self._hist_buffers is None: + capacity = shard_capacity(manager.max_seq_len, self.mapping.tp_size) + if capacity > MAX_SHARD: + return False + device = self._unit_scale.device + tokens_per_block = manager.tokens_per_block + rows = FUSED_A_MAX_TOKENS + sms = num_sms(device.index) + self._hist_buffers = dict( + capacity=capacity, + shard_ctx=torch.zeros(1, dtype=torch.int32, device=device), + shard_block_table=torch.zeros( + (1, -(-capacity // tokens_per_block)), dtype=torch.int32, device=device + ), + schedule=torch.zeros((sms + 1, 2), dtype=torch.int32, device=device), + logits=torch.empty( + (rows, -(-capacity // 256) * 256), dtype=torch.float32, device=device + ), + topk=torch.empty((rows, TOPK), dtype=torch.int32, device=device), + k_fp8=torch.empty((rows, 128), dtype=torch.float8_e4m3fn, device=device), + k_scale=torch.empty((rows, 1), dtype=torch.float32, device=device), + sms=sms, + ) + buffers = self._hist_buffers + hist_shard_meta( + attn_metadata.gen_indexer_kv_lens_cuda_runtime, + attn_metadata.indexer_k_cache_block_offsets[0], + buffers["shard_ctx"], + buffers["shard_block_table"], + self.mapping.tp_size, + self.mapping.tp_rank, + manager.tokens_per_block, + ) + paged_mqa_logits_schedule( + buffers["shard_ctx"], manager.tokens_per_block, buffers["sms"], buffers["schedule"] + ) + return True + + def _attention( + self, + layer, + position_ids, + hidden_states, + attn_metadata, + sharded, + pool_rows: Optional[torch.Tensor] = None, + positions: Optional[torch.Tensor] = None, + ): + """MLA with the DSA indexer over the generation rows: the attention output ahead of + o_proj, [rows, heads * v_head_dim]. + + On an MTP draft step after the first (``sharded``) the rows attend to a given + selection and the indexer does not run: ``pool_rows``, KV pool rows, or + ``positions``, sequence positions the flash decode maps to pool rows itself.""" + mla = layer.self_attn + mqa = mla.mqa + indexer = mqa.indexer + manager = attn_metadata.kv_cache_manager + num_tokens = attn_metadata.num_tokens + seq_len_q = num_tokens // attn_metadata.num_seqs + head_dim = mla.kv_lora_rank + mla.qk_rope_head_dim + mqa._ensure_rope_table_size(attn_metadata.max_seq_len) + cos_sin = mqa.rotary_cos_sin + kv_lens = attn_metadata.kv_lens_cuda_runtime + local_layer = mqa.get_local_layer_idx(attn_metadata) + # Asked per step: the executor re-creates the cache manager after its memory profiling. + kv_pool = manager.get_unique_primary_pool() + reuses_selection = pool_rows is not None or positions is not None + own_indexer = sharded and indexer is not None and not reuses_selection + # A layer that attends to the selection of the indexer layer before it gathers its old KV + # rows before the flash decode's wait. + early_gather = sharded and indexer is None and not reuses_selection + + weight = ( + self._fused_a_indexer_weight[layer.layer_idx] + if own_indexer + else mla.kv_a_proj_with_mqa.weight + ) + if reuses_selection: + l2_prefetch = self._draft.l2_after_fused_a + elif early_gather: + # Also the selection the flash decode reads first. + l2_prefetch = self._l2_after_fused_a[layer.layer_idx] + [ + self._hist_buffers["topk"][:num_tokens] + ] + else: + l2_prefetch = self._l2_after_fused_a[layer.layer_idx] + kv_a = fused_a_gemm(hidden_states, weight, l2_prefetch) + kv_a = kv_a[:num_tokens] + q_a = kv_a[:, : mla.q_lora_rank] + latent_end = mla.q_lora_rank + head_dim + quant_q = hidden_states.new_empty( + (num_tokens, mla.num_heads_tp * head_dim), dtype=torch.uint8 + ) + + def q_b(plan="regular"): + q_b_fp8_rope( + q_a, + self._q_b_absorbed[layer.layer_idx], + quant_q, + kv_lens, + cos_sin.view(-1, mla.qk_rope_head_dim, 2), + self._unit_scale, + mla.q_a_layernorm.variance_epsilon, + head_dim, + mla.qk_rope_head_dim, + seq_len_q, + plan, + ) + + def kv_half(trigger_at_entry=False): + mla_kv_half_append( + kv_a[:, mla.q_lora_rank : latent_end], + mla.kv_a_layernorm.weight, + mla.kv_a_layernorm.variance_epsilon, + cos_sin, + kv_lens, + 0, + seq_len_q, + attn_metadata.kv_cache_block_offsets, + manager.kv_cache_pool_pointers, + manager.kv_cache_pool_mapping, + local_layer, + manager.tokens_per_block, + attn_metadata.max_seq_len, + self._bmm1_scale, + self._bmm2_scale, + mla.softmax_scale, + trigger_at_entry, + ) + + output = hidden_states.new_empty( + (hidden_states.shape[0], mla.num_heads_tp * mla.v_head_dim) + ) + q_heads = quant_q.view(num_tokens, mla.num_heads_tp, head_dim) + if sharded: + if own_indexer: + + def select(): + self._select_rows( + indexer, kv_a, latent_end, q_a, position_ids, attn_metadata, num_tokens, mla + ) + + def kv_half_then_q_b(): + kv_half() + # Queued after the top-k: 18 CTAs leave the SMs to the indexer beside it. + q_b("beside_topk") + + # The indexer on this stream; the KV half and q_b beside it on the aux stream. + maybe_execute_in_parallel( + select, kv_half_then_q_b, mla.ln_events[0], mla.ln_events[1], mla.aux_stream + ) + else: + # With the early gather, q_b at split-K 4 leaves room for the flash decode, which + # the KV half lets launch at its entry. + maybe_execute_in_parallel( + lambda: q_b("early_gather" if early_gather else "regular"), + lambda: kv_half(early_gather), + mla.ln_events[0], + mla.ln_events[1], + mla.aux_stream, + ) + if pool_rows is not None: + sparse_mla_flash_decode( + q_heads, + kv_pool, + pool_rows, + kv_lens, + self._bmm1_scale, + self._bmm2_scale, + mla.v_b_proj, + output, + seq_len_q, + self.index_topk, + ) + return output + page_index_scale, layer_offset = manager.get_primary_pool_page_index_params(local_layer) + sparse_mla_flash_decode( + q_heads, + kv_pool, + self._hist_buffers["topk"][:num_tokens] if positions is None else positions, + kv_lens, + self._bmm1_scale, + self._bmm2_scale, + mla.v_b_proj, + output, + seq_len_q, + self.index_topk, + attn_metadata.block_table[: attn_metadata.num_seqs], + manager.tokens_per_block, + page_index_scale, + layer_offset, + early_gather, + ) + return output + + # The built-in indexer over the whole context, then the selection as cache rows. + maybe_execute_in_parallel(q_b, kv_half, mla.ln_events[0], mla.ln_events[1], mla.aux_stream) + indexer_intermediates = None + if indexer is not None: + qr = mla.q_a_layernorm(q_a) + indexer_intermediates = [ + t[:num_tokens] + for t in indexer.pre_indexer_proj( + qr, hidden_states[:num_tokens], position_ids[..., :num_tokens] + ) + ] + indexer._update_k_cache( + indexer_intermediates[1], indexer_intermediates[2], attn_metadata + ) + indices, _ = mqa.sparse_attn_predict( + hidden_states[:num_tokens], + None, + attn_metadata, + AttentionForwardArgs( + attention_input_type=AttentionInputType.generation_only, + sparse_backend_args=DSABackendForwardArgs( + indexer_intermediates=indexer_intermediates + ), + ), + ) + sparse_mla_flash_decode( + q_heads, + kv_pool, + indices, + kv_lens, + self._bmm1_scale, + self._bmm2_scale, + mla.v_b_proj, + output, + seq_len_q, + self.index_topk, + ) + if indexer is not None: + indexer.maybe_join_prev_topk_copy() + return output + + def _select_rows(self, indexer, kv_a, latent_end, q_a, position_ids, attn_metadata, rows, mla): + """The indexer of a layer that runs one, on this rank's key shard: writes the exact + top-2048 cache positions of every row into the shared selection buffer.""" + buffers = self._hist_buffers + manager = attn_metadata.kv_cache_manager + rope = indexer.rotary_emb + cos_sin = rope.rotary_cos_sin.view(rope.max_positions, -1) + interleave = not rope.is_neox + positions = position_ids.view(-1)[:rows] + index_k = kv_a[:, latent_end : latent_end + indexer.head_dim] + index_weights = kv_a[:, latent_end + indexer.head_dim :] + k_cache = manager.get_indexer_k_cache_buffers(indexer.layer_idx) + + # The K tail on its own stream, beside the Q tail: k_norm, RoPE, FP8 and the K cache. + self._k_tail_fork.record() + with torch.cuda.stream(self._k_tail_stream): + self._k_tail_stream.wait_event(self._k_tail_fork) + indexer_k_tail( + index_k, + positions, + cos_sin, + indexer.k_norm.weight, + indexer.k_norm.bias, + indexer.k_norm.variance_epsilon, + indexer.rope_dim, + interleave, + indexer.scale_fmt == "ue8m0", + buffers["k_fp8"][:rows], + buffers["k_scale"][:rows], + k_cache, + attn_metadata.slot_mapping_fp8[:rows], + attn_metadata.slot_mapping_scale[:rows], + rows, + True, + True, + False, + ) + self._k_tail_done.record() + + q_fp8 = q_a.new_empty((rows, indexer.n_heads, indexer.head_dim), dtype=torch.float8_e4m3fn) + q_scale = q_a.new_empty((rows, indexer.n_heads, 1), dtype=torch.float32) + weights = q_a.new_empty((rows, indexer.n_heads), dtype=torch.float32) + indexer_q_tail( + q_a, + self._wq_b_folded[indexer.layer_idx], + positions, + cos_sin, + index_weights, + q_fp8, + q_scale, + weights, + mla.q_a_layernorm.variance_epsilon, + indexer.n_heads, + indexer.head_dim, + indexer.rope_dim, + interleave, + indexer.scale_fmt == "ue8m0", + indexer.weight_scale_factor, + ) + torch.cuda.current_stream().wait_event(self._k_tail_done) + logits = fp8_paged_mqa_logits( + q_fp8.view(1, rows, indexer.n_heads, indexer.head_dim), + k_cache, + weights, + buffers["shard_ctx"], + buffers["shard_block_table"], + buffers["schedule"], + buffers["logits"][:rows], + buffers["capacity"], + True, + ) + hist_topk( + logits, + attn_metadata.gen_indexer_kv_lens_cuda_runtime, + buffers["topk"][:rows], + self._hist_workspace.comm_buffer, + self._hist_workspace.state, + rows, + self.mapping.tp_size, + ) + + def _post_attention_hop(self, layer, attn_out: torch.Tensor, residual: torch.Tensor): + """o_proj published into the all-reduce workspace, then the hop with the residual add + and post_attention_layernorm: (normed, residual).""" + rows = attn_out.shape[0] + lamport, flags = publish_target(self.mapping, rows, self.hidden_size) + comm_buffer, buffer_flags = hop_workspace(self.mapping, rows, self.hidden_size) + shard = attn_out.new_empty((rows, self.hidden_size)) + o_proj_publish( + attn_out, + layer.self_attn.o_proj.weight, + shard, + lamport, + flags, + self.mapping.tp_size, + self.mapping.tp_rank, + ) + norm = layer.post_attention_layernorm + return mnnvl_hop( + shard, residual, norm.weight, norm.variance_epsilon, comm_buffer, buffer_flags, True + ) + + def _hop(self, norm, partial: torch.Tensor, residual: torch.Tensor): + """The hop of a rank-partial sum with the residual add and `norm`: (normed, residual).""" + comm_buffer, buffer_flags = hop_workspace(self.mapping, partial.shape[0], self.hidden_size) + return mnnvl_hop( + partial, residual, norm.weight, norm.variance_epsilon, comm_buffer, buffer_flags, False + ) + + def _moe(self, layer, hidden_states: torch.Tensor, residual: torch.Tensor): + """The MoE block and the post-MoE hop that finalizes it: (normed, residual).""" + mlp = layer.mlp + backend = getattr(mlp.experts, "backend", mlp.experts) + rows = hidden_states.shape[0] + + self._quant_fork.record() + with torch.cuda.stream(self._quant_stream): + self._quant_stream.wait_event(self._quant_fork) + x_fp4, x_sf = fp4_quantize( + hidden_states, backend.fc31_input_scale, backend.scaling_vector_size, False, False + ) + self._quant_ready.record() + routing = extract_routing_params(backend.routing_method) + + def routed(): + # Unfinalized: [FC2 rows, expert weights, expanded index -> FC2 row]. + return routed_experts( + self._router_logits[:rows], + routing.routing_bias, + x_fp4, + x_sf.flatten().view(torch.float8_e4m3fn), + backend.w3_w1_weight, + backend.w3_w1_weight_scale.view(torch.float8_e4m3fn), + None, + backend.act_alpha, + backend.act_beta, + backend.act_clamp, + backend.w2_weight, + backend.w2_weight_scale.view(torch.float8_e4m3fn), + None, + get_data_or_none(backend, "fc31_scale_c"), + get_data_or_none(backend, "fc31_alpha"), + get_data_or_none(backend, "fc2_alpha"), + backend.num_slots, + routing.top_k, + routing.n_group, + routing.topk_group, + backend.w3_w1_weight.shape[-2] // 2, + backend.slot_start, + backend.expert_size_per_partition, + routing.routed_scaling_factor, + backend.routing_method.routing_method_type, + False, + to_trtllm_gen_act_type(backend.activation_type), + None, + None, + None, + backend.max_num_tokens, + False, + self._quant_ready.cuda_event, + hidden_states, + mlp.gate.weight, + self._router_workspace[layer.layer_idx], + True, + ) + + (fc2_output, expert_weights, permuted_rows), shared_output = maybe_execute_in_parallel( + routed, + lambda: self._shared_expert(mlp, hidden_states), + mlp.event_dict[EventType.Main], + mlp.event_dict[EventType.MoeShared], + mlp.aux_stream, + ) + norm = layer.next_layer_layernorm + comm_buffer, buffer_flags = hop_workspace(self.mapping, rows, self.hidden_size) + return mnnvl_moe_finalize_hop( + fc2_output, + expert_weights, + permuted_rows, + shared_output, + residual, + norm.weight, + norm.variance_epsilon, + comm_buffer, + buffer_flags, + ) + + def _shared_expert(self, mlp, hidden_states: torch.Tensor) -> torch.Tensor: + """This rank's partial sum of the shared expert: two split-K GEMMs around SwiGLU.""" + expert = mlp.shared_experts + rows = hidden_states.shape[0] + gate_up = expert.gate_up_proj.weight + h = low_m_gemm(hidden_states, gate_up, hidden_states.new_empty(rows, gate_up.shape[0])) + h = swiglu(h) + return low_m_gemm( + h, expert.down_proj.weight, hidden_states.new_empty(rows, self.hidden_size) + ) + + @staticmethod + def _builtin_layer(layer, position_ids, hidden_states, residual, attn_metadata): + """One decoder layer on the built-in modules, for a batch the kernels do not serve.""" + dense = isinstance(layer.mlp, GatedMLP) + hidden_states = DecodeTarget._attention_builtin( + layer, + position_ids, + hidden_states, + attn_metadata, + AllReduceParams(enable_allreduce=dense), + ) + if dense: + hidden_states, residual = layer.post_attention_layernorm(hidden_states, residual) + hidden_states = layer.mlp( + hidden_states, final_all_reduce_params=AllReduceParams(enable_allreduce=False) + ) + else: + hidden_states, residual = layer.allreduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=layer.post_attention_layernorm.weight, + eps=layer.post_attention_layernorm.variance_epsilon, + trigger_completion_at_end=False, + ), + ) + hidden_states = layer.mlp( + hidden_states, + None, + all_rank_num_tokens=attn_metadata.all_rank_num_tokens, + final_all_reduce_params=AllReduceParams(enable_allreduce=False), + do_finalize=True, + ) + return layer.allreduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=layer.next_layer_layernorm.weight, + eps=layer.next_layer_layernorm.variance_epsilon, + trigger_completion_at_end=dense, + ), + ) + + @staticmethod + def _attention_builtin(layer, position_ids, hidden_states, attn_metadata, all_reduce_params): + """The built-in modules' MLA with the DSA indexer, for a batch the kernels do not serve.""" + mla = layer.self_attn + num_tokens = attn_metadata.num_tokens + kv_a = mla.kv_a_proj_with_mqa(hidden_states) + q, compressed_kv, k_pe = kv_a.split( + [mla.q_lora_rank, mla.kv_lora_rank, mla.qk_rope_head_dim], -1 + ) + (q, qr), compressed_kv = maybe_execute_in_parallel( + lambda: mla._q_a_layernorm_maybe_fused(q, return_norm_out=True), + lambda: mla.kv_a_layernorm(compressed_kv), + mla.ln_events[0], + mla.ln_events[1], + mla.aux_stream, + ) + latent_cache = torch.concat([compressed_kv, k_pe], dim=-1) + q = mla.q_b_proj(q) + indexer = mla.mqa.indexer + indexer_intermediates = None + if indexer is not None: + indexer_intermediates = [ + t[:num_tokens] for t in indexer.pre_indexer_proj(qr, hidden_states, position_ids) + ] + indexer._update_k_cache( + indexer_intermediates[1], indexer_intermediates[2], attn_metadata + ) + output = hidden_states.new_empty([hidden_states.shape[0], mla.attention_output_hidden_size]) + forward_generation_sparse_attn( + mla, + q[:num_tokens], + compressed_kv[:num_tokens], + k_pe[:num_tokens], + attn_metadata, + output[:num_tokens], + position_ids=position_ids[..., :num_tokens], + latent_cache=latent_cache[:num_tokens], + indexer_intermediates=indexer_intermediates, + ) + if indexer is not None: + indexer.maybe_join_prev_topk_copy() + return mla.o_proj(output, all_reduce_params=all_reduce_params, layer_idx=mla.layer_idx) + + +class Glm52Nvfp4Sm107Tp16Core(DeepseekV3Model): + """The decoder stack, with the per-step choice between the two targets.""" + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + super().__init__(model_config) + self._prefill: Optional[PrefillTarget] = None + self._decode: Optional[DecodeTarget] = None + + @torch.no_grad() + def post_load_weights(self) -> None: + """Build the targets once the weights are final. Reached more than once + (the shell and the loader's module walk), so it only builds once.""" + if self._decode is not None: + return + self._prefill = PrefillTarget(self) + self._decode = DecodeTarget(self) + + def _select_target(self, attn_metadata: AttentionMetadata) -> "PrefillTarget | DecodeTarget": + if self._decode is None: + self.post_load_weights() + return self._prefill if attn_metadata.num_contexts > 0 else self._decode + + def decode_target(self, attn_metadata: AttentionMetadata) -> Optional[DecodeTarget]: + """The decode target when this step runs it, else None.""" + target = self._select_target(attn_metadata) + return target if target is self._decode else None + + def build_draft(self, layer, lm_head, weights: _weights.DraftWeights) -> None: + """Build the decode target's MTP draft-step state (``DecodeTarget.build_draft``).""" + if self._decode is None: + self.post_load_weights() + self._decode.build_draft(layer, lm_head, weights) + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, + **kwargs, + ) -> torch.Tensor: + return self._select_target(attn_metadata).forward( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + **kwargs, + ) + + +class MTPDraftWorker(MTPEagleWorker): + """The one-model MTP-Eagle worker: the draft loop of a one-request generation step runs + on the decode target's kernels (``DecodeTarget``'s draft steps); every other step runs the + built-in loop. + + The loop keeps the built-in loop's attention-metadata handling: the MTP index-share flags, + the switch to one row per request after draft step 0, and the cleanup. The draft-step + kernels advance the sequence and KV lengths in place, and the draft tokens and the + next-new-tokens rows come from them. + """ + + def __init__( + self, + spec_config, + model_config: ModelConfig, + use_separate_draft_kv_cache: bool, + *, + mapping, + decode_target, + ) -> None: + super().__init__(spec_config, model_config, use_separate_draft_kv_cache, mapping=mapping) + # The core's per-step lookup of the decode target (a bound method rather than the + # core itself, which belongs to the shell and is no submodule of this worker). + self._decode_target = decode_target + # The draft tokens of the last draft loop on the decode target and the next-new-tokens + # rows its last step wrote, for _prepare_next_new_tokens. + self._drafted_rows: Optional[Tuple[torch.Tensor, torch.Tensor]] = None + + def prepare_1st_drafter_inputs( + self, + input_ids, + position_ids, + hidden_states, + accepted_tokens, + attn_metadata, + spec_metadata, + draft_model, + ): + """Without context requests draft step 0's input ids are the accepted-token rows as + they are, flattened, rather than a copy behind an empty context part.""" + if not self.is_mtp_eagle or attn_metadata.num_contexts > 0: + return super().prepare_1st_drafter_inputs( + input_ids, + position_ids, + hidden_states, + accepted_tokens, + attn_metadata, + spec_metadata, + draft_model, + ) + return { + "input_ids": accepted_tokens[:, : spec_metadata.runtime_draft_len + 1].flatten(), + "position_ids": position_ids, + "hidden_states": hidden_states, + "attn_metadata": attn_metadata, + "spec_metadata": spec_metadata, + } + + def _prepare_next_new_tokens( + self, + accepted_tokens, + next_draft_tokens, + batch_indices_cuda, + batch_size, + num_accepted_tokens, + ): + """The rows the last draft-step kernel wrote when these draft tokens are that loop's.""" + drafted, self._drafted_rows = self._drafted_rows, None + if drafted is not None and drafted[0] is next_draft_tokens: + return drafted[1] + return super()._prepare_next_new_tokens( + accepted_tokens, + next_draft_tokens, + batch_indices_cuda, + batch_size, + num_accepted_tokens, + ) + + def _sample_and_accept_draft_tokens_base( + self, + logits, + draft_tokens, + num_contexts, + batch_size, + spec_metadata, + ): + """Strict acceptance and the forced accepted-token count in one kernel + (``glm_verify_head_accept``) for a step of generation requests only, with the built-in + target-token sampling; ragged verification and context requests take the built-in + chain. The forced count follows ``_apply_force_accepted_tokens``: none during eager + CUDA-graph warm-up, the integer part plus one, and with a fractional part the same pool + draw and counter step.""" + if spec_metadata is None or spec_metadata.is_ragged_verify or num_contexts != 0: + return super()._sample_and_accept_draft_tokens_base( + logits, draft_tokens, num_contexts, batch_size, spec_metadata + ) + if logits.dim() == 1: + logits = logits.unsqueeze(0) + tokens_per_request = draft_tokens.shape[-1] + 1 + target_tokens = self._sample_tokens_for_batch(logits, spec_metadata, 0, batch_size) + mode, base, frac = 0, 0, 0.0 + warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + if self.force_num_accepted_tokens != 0.0 and not warmup: + int_part = int(self.force_num_accepted_tokens) + frac = self.force_num_accepted_tokens - int_part + base = min(int_part + 1, tokens_per_request) + mode = 2 if frac > 0.0 and base < tokens_per_request else 1 + pool = counter = None + if mode == 2: + self._ensure_force_accept_rng_state(logits.device) + pool, counter = self._force_accept_rng_pool, self._force_accept_rng_counter + accepted_tokens = torch.empty( + (batch_size, tokens_per_request), dtype=torch.int, device=logits.device + ) + num_accepted_tokens = glm_verify_head_accept( + target_tokens, + draft_tokens, + 0, + accepted_tokens, + pool, + counter, + mode, + base, + tokens_per_request, + frac, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + ) + return accepted_tokens, num_accepted_tokens + + def _forward_linear_draft_loop( + self, + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + ): + self._drafted_rows = None + decode = self._drafting_target( + inputs, + attn_metadata, + spec_metadata, + draft_kv_cache_manager, + num_contexts, + batch_size, + num_accepted_tokens, + ) + if decode is None: + return super()._forward_linear_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + ) + return self._draft_on_target( + decode, + inputs, + attn_metadata, + spec_metadata, + draft_kv_cache_manager, + batch_size, + num_accepted_tokens, + ) + + def _drafting_target( + self, + inputs, + attn_metadata, + spec_metadata, + draft_kv_cache_manager, + num_contexts, + batch_size, + num_accepted_tokens, + ) -> Optional[DecodeTarget]: + """The decode target when its kernels draft this step, else None. + + They draft a generation step of one request whose draft layer shares the target's KV + pool, through the configured draft length with no draft-length schedule (the draft + steps after the first skip the MTP layer's indexer K write: the next step's draft + step 0 writes those positions before it reads them), greedily (argmax drafts), with + nothing that rewrites the drafts (guided decoding, SA, a draft vocabulary map), and + when the decode target's kernels serve the rows. + """ + draft_len = spec_metadata.runtime_draft_len + rows = batch_size * (draft_len + 1) + input_ids = inputs["input_ids"] + position_ids = inputs["position_ids"] + if not ( + self.is_mtp_eagle + and num_contexts == 0 + and batch_size == 1 + and draft_kv_cache_manager is None + and draft_len == self.max_draft_len + and getattr(self.spec_config, "draft_len_schedule", None) is None + and self.guided_decoder is None + and self.sa_enhancer is None + and self._d2t is None + and not spec_metadata.wants_advanced_draft_sampling + and not spec_metadata.is_ragged_verify + and attn_metadata.num_tokens == rows + and inputs["hidden_states"].shape[0] == rows + and input_ids.dtype == torch.int32 + and tuple(input_ids.shape) == (rows,) + and input_ids.is_contiguous() + and position_ids.dtype in (torch.int32, torch.int64) + and tuple(position_ids.shape) == (rows,) + and position_ids.is_contiguous() + and num_accepted_tokens.dtype == torch.int32 + and num_accepted_tokens.is_contiguous() + and spec_metadata.batch_indices_cuda.dtype == torch.int32 + ): + return None + decode = self._decode_target(attn_metadata) + if decode is None or not decode.drafts(rows): + return None + return decode + + def _draft_on_target( + self, + decode: DecodeTarget, + inputs, + attn_metadata, + spec_metadata, + draft_kv_cache_manager, + batch_size, + num_accepted_tokens, + ) -> torch.Tensor: + """The draft loop on the decode target's kernels: the draft tokens [batch, draft_len]. + + Each step's last kernel writes the next step's sequence and KV lengths in place. The + steps after the first run no indexer and read nothing else of the DSA decode metadata, + so the built-in loop's per-step rebuild of it (``update_for_spec_dec``) is left out; + with a context short enough that the metadata skips the indexer, only the next step's + dense selection is refreshed from the new lengths. The host side of draft step 0's + switch to one row per request comes before that. + """ + draft_len = spec_metadata.runtime_draft_len + # The accepted tokens' rows: draft step 0's input ids (prepare_1st_drafter_inputs). + accepted_tokens = inputs["input_ids"].view(batch_size, draft_len + 1) + batch_indices = spec_metadata.batch_indices_cuda[:batch_size] + num_accepted = num_accepted_tokens[:batch_size] + draft_tokens = accepted_tokens.new_empty((batch_size, draft_len)) + positions = inputs["position_ids"] + eh_proj_output = gather_ids = pool_rows = next_new_tokens = None + uses_mtp_index_share = isinstance(attn_metadata, MTPIndexShareMetadata) + with contextlib.ExitStack() as draft_scope: + if uses_mtp_index_share: + attn_metadata.set_in_mtp_draft_loop(True) + attn_metadata.set_mtp_num_accepted(num_accepted_tokens) + draft_scope.callback(_reset_mtp_index_share, attn_metadata) + attn_metadata = draft_scope.enter_context( + self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager) + ) + # The draft loop's copy of the sequence lengths (prepare_for_spec_dec) and the KV + # lengths, which the draft-step kernels advance. + seq_lens = attn_metadata.seq_lens_cuda[:batch_size] + kv_lens = attn_metadata.kv_lens_cuda[:batch_size] + for i in range(draft_len): + if uses_mtp_index_share: + attn_metadata.set_skip_topk(i > 0) + if i == 0: + hidden_states, gather_ids, pool_rows = decode.draft_first_step( + attn_metadata, + inputs["input_ids"], + positions, + inputs["hidden_states"], + batch_indices, + num_accepted, + draft_len, + ) + else: + hidden_states = decode.draft_next_step( + attn_metadata, eh_proj_output, pool_rows, join_selection=i == 1 + ) + last = i == draft_len - 1 + logits, hidden_half = decode.draft_head(hidden_states, next_step=not last) + rewind_accepted = num_accepted if i == 0 else None + if i == 0: + # One row per request from here on, on the host (no context requests to + # turn into generation ones); the kernel below fills seq_lens. + attn_metadata._seq_lens[:batch_size].fill_(1) + attn_metadata.on_update() + if last: + next_new_tokens = decode.draft_sample_glue( + logits, + draft_tokens, + i, + seq_lens, + kv_lens, + rewind_accepted, + accepted_tokens, + batch_indices, + num_accepted, + ) + else: + eh_proj_output, positions = decode.draft_ehproj_table( + logits, + hidden_half, + draft_tokens, + i, + positions, + gather_ids if i == 0 else None, + seq_lens, + kv_lens, + rewind_accepted, + ) + if not last and attn_metadata.skip_indexer_for_gen_reqs: + attn_metadata.prepare_dense_topk_indices( + attn_metadata.kv_lens_cuda, device=True + ) + if i == 0: + attn_metadata.use_spec_decoding = False + self._drafted_rows = (draft_tokens, next_new_tokens) + return draft_tokens + + +class ModelingV2Glm52Nvfp4Sm107Tp16( + SpecDecOneEngineForCausalLM[Glm52Nvfp4Sm107Tp16Core, PretrainedConfig] +): + """GLM-5.2-NVFP4, sm107, tp16.""" + + @classmethod + def get_preferred_kv_cache_manager_version(cls, pretrained_config: Any = None) -> Literal["V2"]: + """The KV cache manager the context server's built-in implementation prefers.""" + return "V2" + + @classmethod + def get_preferred_transceiver_runtime( + cls, pretrained_config: Any = None + ) -> Optional[Literal["CPP", "PYTHON"]]: + """The transceiver runtime for the masked DSA indexer K-cache transfer.""" + return "PYTHON" + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + mapping = model_config.mapping + assert torch.cuda.get_device_capability() == _COMPUTE_CAPABILITY, ( + f"this target runs on sm107, not {torch.cuda.get_device_capability()}" + ) + assert ( + mapping.world_size == _TP + and mapping.tp_size == _TP + and mapping.moe_ep_size == 1 + and not mapping.enable_attention_dp + and mapping.pp_size == 1 + ), f"this target is tp16 without attention DP or expert parallelism, not {mapping}" + model_config = _weights.normalize_model_config(model_config) + super().__init__(model=Glm52Nvfp4Sm107Tp16Core(model_config), model_config=model_config) + if ( + model_config.spec_config is not None + and model_config.spec_config.spec_dec_mode.is_mtp_one_model() + ): + self.num_hidden_layers = self.config.num_hidden_layers + self.model.layers.extend(self.draft_model.mtp_layers) + if type(self.spec_worker) is MTPEagleWorker: + # The target's worker in place of the built-in one, with the same draft KV cache + # choice and draft model. + worker = MTPDraftWorker( + self.spec_config, + model_config, + self.use_separate_draft_kv_cache, + mapping=model_config.mapping, + decode_target=self.model.decode_target, + ) + worker.set_draft_model(self.draft_model) + builtin = self.spec_worker + self.epilogue[:] = [worker if m is builtin else m for m in self.epilogue] + self.spec_worker = worker + # The MTP layer's checkpoint tensors the draft steps derive operands from, held from + # load_weights to post_load_weights. + self._draft_weights: Optional[_weights.DraftWeights] = None + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: torch.IntTensor = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + return_context_logits: bool = False, + spec_metadata: Optional[SpecMetadata] = None, + resource_manager=None, + **kwargs, + ) -> torch.Tensor: + decode = self.model.decode_target(attn_metadata) + if decode is None or self.spec_worker is None: + return super().forward( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + return_context_logits=return_context_logits, + spec_metadata=spec_metadata, + resource_manager=resource_manager, + **kwargs, + ) + # A generation step with speculative decoding: the verify logits on the decode + # target's kernels where they serve the rows, then the spec worker. + hidden_states = self.model( + input_ids=input_ids, + attn_metadata=attn_metadata, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + **kwargs, + ) + if attn_metadata.padded_num_tokens is not None: + num_tokens = attn_metadata.num_tokens + hidden_states = hidden_states[:num_tokens] + if input_ids is not None: + input_ids = input_ids[:num_tokens] + position_ids = _slice_spec_position_ids(position_ids, num_tokens) + rows = hidden_states[spec_metadata.gather_ids] + logits = decode.logits(self.lm_head, rows) + if logits is None: + logits = self.logits_processor.forward(rows, self.lm_head, attn_metadata, True) + return self.spec_worker( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + logits=logits, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=self.draft_model, + resource_manager=resource_manager, + ) + + def load_weights(self, weights, *args, **kwargs): + self._draft_weights = _weights.load_weights( + self, weights, isinstance(self.spec_worker, MTPDraftWorker) + ) + + def setup_aliases(self) -> None: + DeepseekV3ForCausalLM.setup_aliases(self) + + def post_load_weights(self): + super().post_load_weights() + self.model.post_load_weights() + draft_weights, self._draft_weights = self._draft_weights, None + if draft_weights is not None: + self.model.build_draft(self.draft_model.mtp_layers[0], self.lm_head, draft_weights) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/__init__.py new file mode 100644 index 000000000000..f8bbe8b7c789 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Entries only this target calls, one module per op under a category package.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/mla_kv_half_append.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/mla_kv_half_append.py new file mode 100644 index 000000000000..948cc51554d8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/mla_kv_half_append.py @@ -0,0 +1,292 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The KV half of an MLA generation step: kv_a RMSNorm, K RoPE and the FP8 paged-cache append.""" + +from __future__ import annotations + +import math + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The latent row: kv_lora_rank 512 normalized columns followed by qk_rope_head_dim 64 K-RoPE +#: columns. +KV_LORA_RANK = 512 +ROPE_DIM = 64 +#: Where the latent starts in fused-A's output row (after q_lora_rank 2048 q_a columns). +LATENT_OFFSET = 2048 +#: The E4M3 range the kernel saturates to. +_E4M3_MAX = 448.0 +#: log2(e) as the kernel's fp32 constant. +_LOG2E = 1.4426950408889634 + +#: GLM-5.2's pool page, kv_a_layernorm epsilon, 1 / sqrt(qk_nope_head_dim + qk_rope_head_dim) +#: and rope_theta. +_GLM = dict( + tokens_per_block=64, + kv_norm_eps=1e-5, + softmax_scale=1.0 / math.sqrt(256.0), + rope_theta=8_000_000.0, +) + + +class GlmMlaKvHalfAppend(OpWrapper): + """One generation step's MLA KV half over a paged FP8 latent cache. + + Wraps ``torch.ops.trtllm.glm_mla_kv_half_append``. Row ``t`` of ``latent`` is token + ``t % seq_len_q`` of generation sequence ``s = t // seq_len_q`` and sits at position + ``p = sequence_length[num_contexts + s] - seq_len_q + t % seq_len_q``. Its cache row is:: + + e4m3(rmsnorm(latent[t, :512], kv_norm_weight, kv_norm_eps)) # x * (rstd * w), fp32 + e4m3(gptj_rope(latent[t, 512:576], position p)) # pairs (2i, 2i + 1), fp32 + + saturated to +-448 at unit scale, written to layer ``layer_idx`` of the paged pool at + position ``p`` through the same block offsets, pool pointers and pool mapping as + ``mla_rope_generation``. A negative position is skipped. The call also fills + ``mla_bmm1_scale = [softmax_scale, softmax_scale * log2(e)]`` and ``mla_bmm2_scale = [1]``, + the scales the generation attention reads. It returns nothing. + + Values reach E4M3 straight from fp32: unlike ``mla_rope_generation``'s KV half, neither + the normalized latent nor the rotated K is rounded to BF16 first. + + Fusion boundary: the RMSNorm of the latent, the RoPE of K, the quantization, the append + and the scales are inside. kv_a_proj, the query and the attention are not. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="verify and MTP step 0: 8 rows of fused-A's 2624-column output across a page edge", + spec=dict( + seq_len_q=8, + cached=(4_090,), + latent_width=2624, + num_layers=4, + layer_idx=1, + trigger_at_entry=True, + **_GLM, + ), + ), + Cell( + why="full-indexer layers: the latent inside fused-A's 2784-column output (wk|wp rows)", + spec=dict( + seq_len_q=8, + cached=(4_090,), + latent_width=2784, + num_layers=4, + layer_idx=0, + trigger_at_entry=True, + **_GLM, + ), + ), + Cell( + why="MTP draft steps 1-6: one row at the MTP layer, the last layer of the pool", + spec=dict( + seq_len_q=1, + cached=(4_100,), + latent_width=2624, + num_layers=4, + layer_idx=3, + trigger_at_entry=True, + **_GLM, + ), + ), + Cell( + why="two sequences (max_batch_size 2): per-sequence offsets and positions, end trigger", + spec=dict( + seq_len_q=8, + cached=(700, 4_090), + latent_width=2624, + num_layers=4, + layer_idx=2, + trigger_at_entry=False, + **_GLM, + ), + ), + ) + + note = """ + Metadata consumed: none -- every input is an argument. What the target passes: + - `latent`: fused-A's output columns [2048, 2624) -- the raw kv_a_proj row, not + normalized; `kv_norm_weight` / `kv_norm_eps` the kv_a_layernorm weight and epsilon. + - `rotary_cos_sin`: the MLA rope table (`RopeParams(..., duplicate_data=True)`, GPT-J + pairs), covering every position the cache can hold (`max_seq_len`). + - `sequence_length`, `num_contexts`, `kv_cache_block_offsets`: the attention metadata's + `kv_lens_cuda_runtime`, `num_contexts` and `kv_cache_block_offsets`; + `host_kv_cache_pool_pointers` / `host_kv_cache_pool_mapping` the cache manager's + `kv_cache_pool_pointers` / `kv_cache_pool_mapping`; `layer_idx` the local layer; + `attention_window_size` the metadata's `max_seq_len`. + - `mla_bmm1_scale` ([2] fp32) and `mla_bmm2_scale` ([1] fp32): buffers allocated once, + read by `glm_sparse_mla_flash_decode`. + + The cache carries no scaling factor: GLM-5.2's FP8 cache is written and read at 1.0, as + the PyTorch MLA path writes it on the context server. The pool must hold plain 576-byte + E4M3 rows (`kv_cache_dtype: fp8`); the inline-scale `fp8_ds_mla` layout has 656-byte rows, + which the paged-cache arguments do not reveal, so the target asserts the KV dtype at + construction. + + PDL: the kernel waits on its grid dependency before any global read. With + `trigger_at_entry` its dependents may launch at its entry (the flash decode that gathers + earlier steps' rows before its wait), otherwise after its stores; every dependent must + wait on its own grid dependency before reading the appended rows or the scales. + + Not certified: a mixed batch (`num_contexts` > 0), more than two sequences, a ragged batch, + pages other than 64 tokens. Rewriting a prepared step rewrites the same slots. + """ + + def __call__( + self, + latent: torch.Tensor, + kv_norm_weight: torch.Tensor, + kv_norm_eps: float, + rotary_cos_sin: torch.Tensor, + sequence_length: torch.Tensor, + num_contexts: int, + seq_len_q: int, + kv_cache_block_offsets: torch.Tensor, + host_kv_cache_pool_pointers: torch.Tensor, + host_kv_cache_pool_mapping: torch.Tensor, + layer_idx: int, + tokens_per_block: int, + attention_window_size: int, + mla_bmm1_scale: torch.Tensor, + mla_bmm2_scale: torch.Tensor, + softmax_scale: float, + trigger_at_entry: bool = False, + ) -> None: + torch.ops.trtllm.glm_mla_kv_half_append( + latent, + kv_norm_weight, + kv_norm_eps, + rotary_cos_sin, + sequence_length, + num_contexts, + seq_len_q, + kv_cache_block_offsets, + host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping, + layer_idx, + tokens_per_block, + attention_window_size, + mla_bmm1_scale, + mla_bmm2_scale, + softmax_scale, + trigger_at_entry, + ) + + def reference( + self, + latent: torch.Tensor, + kv_norm_weight: torch.Tensor, + kv_norm_eps: float, + rotary_cos_sin: torch.Tensor, + sequence_length: torch.Tensor, + num_contexts: int, + seq_len_q: int, + kv_cache_block_offsets: torch.Tensor, + host_kv_cache_pool_pointers: torch.Tensor, + host_kv_cache_pool_mapping: torch.Tensor, + layer_idx: int, + tokens_per_block: int, + attention_window_size: int, + mla_bmm1_scale: torch.Tensor, + mla_bmm2_scale: torch.Tensor, + softmax_scale: float, + trigger_at_entry: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """``(rows, positions, bmm1, bmm2)``: what the op is supposed to write, and where. + + ``rows`` [T, 576] E4M3 is the cache row of each latent row, ``positions`` [T] int64 its + position in its sequence. The paged-cache arguments (block offsets, pool pointers and + mapping, layer, page size, window) only say where a row lands, and the test reads it + back there; ``mla_bmm1_scale`` / ``mla_bmm2_scale`` are where the scales land and + ``trigger_at_entry`` when the dependents launch, so all of them are ignored here. + + fp32 throughout, as the kernel: the RMSNorm multiplies ``x * (rstd * w)``, the RoPE + pairs consecutive columns (GPT-J), and each value is saturated and rounded to E4M3 + once. + """ + num_tokens = latent.shape[0] + device = latent.device + token = torch.arange(num_tokens, device=device) + seq = token // seq_len_q + positions = sequence_length[num_contexts + seq].long() - seq_len_q + token % seq_len_q + + x = latent[:, :KV_LORA_RANK].float() + rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + kv_norm_eps) + nope = x * (rstd * kv_norm_weight.float()) + + pairs = ( + latent[:, KV_LORA_RANK : KV_LORA_RANK + ROPE_DIM] + .float() + .reshape(num_tokens, ROPE_DIM // 2, 2) + ) + table = rotary_cos_sin.reshape(-1, ROPE_DIM, 2)[positions.clamp(min=0), : ROPE_DIM // 2] + cos, sin = table[..., 0], table[..., 1] + roped = torch.stack( + (cos * pairs[..., 0] - sin * pairs[..., 1], cos * pairs[..., 1] + sin * pairs[..., 0]), + dim=-1, + ).reshape(num_tokens, ROPE_DIM) + + rows = torch.cat((nope, roped), dim=-1).clamp(-_E4M3_MAX, _E4M3_MAX).to(torch.float8_e4m3fn) + scale = torch.tensor(softmax_scale, dtype=torch.float32, device=device) + log2e = torch.tensor(_LOG2E, dtype=torch.float32, device=device) + bmm1 = torch.stack((scale, scale * log2e)) + bmm2 = torch.ones(1, dtype=torch.float32, device=device) + return rows, positions, bmm1, bmm2 + + def is_valid( + self, + latent: torch.Tensor, + kv_norm_weight: torch.Tensor, + kv_norm_eps: float, + rotary_cos_sin: torch.Tensor, + sequence_length: torch.Tensor, + num_contexts: int, + seq_len_q: int, + kv_cache_block_offsets: torch.Tensor, + host_kv_cache_pool_pointers: torch.Tensor, + host_kv_cache_pool_mapping: torch.Tensor, + layer_idx: int, + tokens_per_block: int, + attention_window_size: int, + mla_bmm1_scale: torch.Tensor, + mla_bmm2_scale: torch.Tensor, + softmax_scale: float, + trigger_at_entry: bool = False, + ) -> None: + """The lengths the op derives positions from without checking. + + Only ``latent``'s row count, ``sequence_length``, ``num_contexts``, ``seq_len_q``, + ``rotary_cos_sin``'s size, ``kv_cache_block_offsets``' width and ``tokens_per_block`` + are read. + """ + num_generations = latent.shape[0] // seq_len_q + if num_generations == 0: + return + lens = sequence_length[num_contexts : num_contexts + num_generations].long() + # A sequence shorter than its step leaves rows at negative positions, which the kernel + # skips: those tokens never reach the cache. + assert bool((lens >= seq_len_q).all()), ( + "sequence_length must count this step's seq_len_q tokens; a shorter sequence " + "leaves its rows unwritten" + ) + last_position = int(lens.max()) - 1 + # The table is read at the position without a bound. + num_positions = rotary_cos_sin.numel() // (2 * ROPE_DIM) + assert last_position < num_positions, ( + f"rotary_cos_sin covers {num_positions} positions; position {last_position} " + "reads past it" + ) + # A block past the offsets' width reads another sequence's offsets and writes its page. + width = kv_cache_block_offsets.shape[-1] + assert last_position // tokens_per_block < width, ( + f"kv_cache_block_offsets has {width} blocks per sequence; position {last_position} " + "needs more" + ) + + +mla_kv_half_append = GlmMlaKvHalfAppend() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/sparse_mla_flash_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/sparse_mla_flash_decode.py new file mode 100644 index 000000000000..fea70f687292 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/sparse_mla_flash_decode.py @@ -0,0 +1,363 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DSA sparse MLA generation attention at 4 query heads per rank, with the v_b BMM folded in.""" + +from __future__ import annotations + +import math +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +#: Query heads per rank: GLM-5.2's 64 MLA heads over 16 ranks. +NUM_HEADS = 4 +#: The latent row: kv_lora_rank 512 followed by qk_rope_head_dim 64. +HEAD_DIM_QK = 576 +KV_LORA_RANK = 512 +V_HEAD_DIM = 256 +#: The most selected tokens per query token the kernel reads. +MAX_TOP_K = 2048 +#: The kernel splits the top-k slots over 16 CTAs of 128 slots; each split exponentiates +#: against its own max. +_SPLITS = 16 +_SLOTS_PER_SPLIT = MAX_TOP_K // _SPLITS +#: P is quantized to E4M3 after this scale, as the TRTLLM-Gen generation kernels feed their +#: second GEMM. +_P_SCALE = 448.0 + +#: GLM-5.2's KV pool at TP16: 79 attention layers per rank (78 decoder layers and the MTP +#: layer, at local layer 78), one FP8 page of 64 tokens per layer and pool slot. +_GLM_POOL = dict(tokens_per_block=64, num_layers=79, top_k=MAX_TOP_K) +#: 1 / sqrt(qk_nope_head_dim + qk_rope_head_dim) = 1 / sqrt(192 + 64). +GLM_SOFTMAX_SCALE = 1.0 / math.sqrt(256.0) + + +def _cell(**spec) -> dict: + """A cell spec over GLM-5.2's pool geometry. + + ``kv_lens`` holds each sequence's KV length with this step's ``seq_len_q`` tokens; ``remap`` + passes positions and the block table (else global pool rows); ``layer_offset`` is the pool + layer; ``holes`` leaves masked entries inside the rows' valid range. + """ + return {**_GLM_POOL, "gather_before_wait": False, "holes": False, **spec} + + +class GlmSparseMlaFlashDecode(OpWrapper): + """Sparse MLA flash decode over an FP8 paged latent pool, v_b folded into the output. + + Wraps ``torch.ops.trtllm.glm_sparse_mla_flash_decode``. Query token ``t`` (row ``t`` of + ``q``) belongs to sequence ``s = t // seq_len_q`` and attends to the first + ``n_t = min(seq_lens_kv[s] - (seq_len_q - 1 - t % seq_len_q), top_k)`` slots of + ``indices[t]``; a negative entry is masked. With ``block_table`` a slot holds a position + ``p`` of the sequence, read at pool row:: + + block_table[s, p // tokens_per_block] * page_index_scale * tokens_per_block + + p % tokens_per_block + layer_offset * tokens_per_block + + (``convert_req_index_to_global``'s rows; a block past the table's width or a negative + entry masks the slot). Without it, a slot holds the pool row. Every head ``h`` then gets + the softmax over the selected rows of ``(q[t, h] . row) * mla_bmm1_scale[1]`` (base 2), + applied to the rows' first 512 values and scaled by ``mla_bmm2_scale[0]``, rounded to BF16, + and multiplied by ``v_b_proj[h]``: ``output[t, h * 256 + v]``. A row with nothing selected + returns zeros. + + The probabilities enter the second product as E4M3: each split of 128 slots exponentiates + against its own max, and ``exp2(s - max) * 448`` is rounded to E4M3, while the + normalizing sums add the unrounded values. ``reference`` states that rounding. + + Fusion boundary: the index remap, the gather of the selected rows, the softmax, P.V, the + 16-way merge and the v_b product are inside. Selecting the tokens, the KV append, the + query's projection, RoPE and quantization, and the scales are not. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="verify, full-indexer layer: 8 rows, the layer's own top-k remapped in-kernel", + spec=_cell(seq_len_q=8, kv_lens=(20_011,), remap=True, layer_offset=6), + ), + Cell( + why="verify, shared-indexer layer: the leader's top-k, old rows gathered before wait", + spec=_cell( + seq_len_q=8, kv_lens=(20_011,), remap=True, layer_offset=7, gather_before_wait=True + ), + ), + Cell( + why="MTP draft step 0: 8 rows at the MTP layer, the last layer of the pool", + spec=_cell(seq_len_q=8, kv_lens=(20_018,), remap=True, layer_offset=78), + ), + Cell( + why="MTP draft steps 1-6: one row on global pool rows (reused top-k remapped aside)", + spec=_cell(seq_len_q=1, kv_lens=(20_019,), remap=False, layer_offset=78), + ), + Cell( + why="short context: fewer positions than the top-k, -1 padding and masked holes", + spec=_cell( + seq_len_q=8, + kv_lens=(1_500,), + remap=True, + layer_offset=4, + gather_before_wait=True, + holes=True, + ), + ), + Cell( + why="two sequences: per-sequence block-table rows and KV lengths (max_batch_size 2)", + spec=_cell( + seq_len_q=8, + kv_lens=(9_000, 17_003), + remap=True, + layer_offset=11, + gather_before_wait=True, + ), + ), + ) + + note = """ + Metadata consumed: none -- every input is an argument. What the target passes: + - `q`: the [T, 4, 576] E4M3 query that q_b's fused epilogue writes (absorbed q_nope, + roped q_pe); `seq_len_q` the query tokens per sequence (8 on verify and MTP draft + step 0, 1 on draft steps 1-6). + - `kv_pool`: the DSA cache manager's `get_unique_primary_pool()`, the base that + `host_kv_cache_pool_pointers[pool, 0]` names and `convert_req_index_to_global`'s + rows count from. + - With `block_table` (layers that run the indexer, and the shared layers after one): + `indices` the [T, 2048] local top-k positions (the full layer's own; on a shared + layer its full layer's), `block_table` the DSA metadata's `block_table[num_contexts: + num_seqs]` (pool slots, -1 past each sequence), `(page_index_scale, layer_offset)` + `get_primary_pool_page_index_params(local_layer)`, `tokens_per_block` the manager's. + Without it (draft steps 1-6): `indices` are global pool rows. + - `seq_lens_kv`: the generation sequences' KV lengths with this step's tokens + (`kv_lens_cuda_runtime[num_contexts:num_seqs]`). + - `mla_bmm1_scale` / `mla_bmm2_scale`: the buffers `glm_mla_kv_half_append` fills. + + Masked slots load pool row 0 and multiply it by P = 0: a NaN in that row reaches the + output. Outputs are bitwise deterministic. `output` rows past `q.size(0)` are untouched. + + PDL: the kernel waits on its grid dependency before any global read, except with + `gather_before_wait`: then the indices, the block table, the KV lengths and the pool rows + of positions below `seq_lens_kv - seq_len_q` are read before the wait, so their writers + must have completed before this grid launches (step-start host copies, earlier steps, + and an earlier layer's top-k that some kernel ahead in the stream waited on). A layer + whose own top-k writes `indices` must not set it. Dependents may launch once the partial + outputs are sent; they must wait on their grid dependency before reading `output`. + + The test also replays the shared-layer cell from a captured CUDA graph (the op allocates + nothing and syncs nothing) and runs both ops on a DSACacheManagerV2 pool against + `convert_req_index_to_global`. + + Not certified: `top_k` below 2048, more than two sequences, a ragged batch (present it + token-major: `seq_len_q` 1 with per-row lengths and tables). + """ + + def __call__( + self, + q: torch.Tensor, + kv_pool: torch.Tensor, + indices: torch.Tensor, + seq_lens_kv: torch.Tensor, + mla_bmm1_scale: torch.Tensor, + mla_bmm2_scale: torch.Tensor, + v_b_proj: torch.Tensor, + output: torch.Tensor, + seq_len_q: int, + top_k: int, + block_table: Optional[torch.Tensor] = None, + tokens_per_block: int = 0, + page_index_scale: int = 0, + layer_offset: int = 0, + gather_before_wait: bool = False, + ) -> torch.Tensor: + torch.ops.trtllm.glm_sparse_mla_flash_decode( + q, + kv_pool, + indices, + seq_lens_kv, + mla_bmm1_scale, + mla_bmm2_scale, + v_b_proj, + output, + seq_len_q, + top_k, + block_table, + tokens_per_block, + page_index_scale, + layer_offset, + gather_before_wait, + ) + return output + + def reference( + self, + q: torch.Tensor, + kv_pool: torch.Tensor, + indices: torch.Tensor, + seq_lens_kv: torch.Tensor, + mla_bmm1_scale: torch.Tensor, + mla_bmm2_scale: torch.Tensor, + v_b_proj: torch.Tensor, + output: torch.Tensor, + seq_len_q: int, + top_k: int, + block_table: Optional[torch.Tensor] = None, + tokens_per_block: int = 0, + page_index_scale: int = 0, + layer_offset: int = 0, + gather_before_wait: bool = False, + ) -> torch.Tensor: + """The [T, 1024] BF16 rows the op writes into ``output[:T]``. + + fp32 throughout, with the kernel's three roundings stated: P to E4M3 at scale 448 per + 128-slot split (against the split's max, the sums over the unrounded P), the merged + attention output to BF16 before the v_b product, and the product to BF16. ``output`` + and ``gather_before_wait`` are ignored: one is where the op writes, the other when it + reads. + """ + num_tokens = q.shape[0] + device = q.device + pool_rows = kv_pool.view(torch.uint8).reshape(-1, HEAD_DIM_QK) + + slots = torch.full((num_tokens, MAX_TOP_K), -1, dtype=torch.int64, device=device) + slots[:, :top_k] = indices[:, :top_k].long() + token = torch.arange(num_tokens, device=device) + seq = token // seq_len_q + num_valid = torch.clamp( + seq_lens_kv[seq].long() - (seq_len_q - 1 - token % seq_len_q), max=top_k + ) + if block_table is not None: + position = slots.clamp(min=0) + block = position // tokens_per_block + width = block_table.shape[1] + page = block_table[seq[:, None], block.clamp(max=width - 1)].long() + row = ( + page * (page_index_scale * tokens_per_block) + + position % tokens_per_block + + layer_offset * tokens_per_block + ) + resident = (slots >= 0) & (block < width) & (page >= 0) + else: + row = slots + resident = slots >= 0 + valid = resident & (torch.arange(MAX_TOP_K, device=device)[None, :] < num_valid[:, None]) + row = torch.where(valid, row, torch.zeros_like(row)) + + keys = pool_rows[row].view(torch.float8_e4m3fn).float() + queries = q.view(torch.float8_e4m3fn).float() + with true_fp32_matmul(): + scores = torch.einsum("thd,tjd->thj", queries, keys) * mla_bmm1_scale[1].float() + scores = scores.masked_fill(~valid[:, None, :], float("-inf")) + scores = scores.view(num_tokens, NUM_HEADS, _SPLITS, _SLOTS_PER_SPLIT) + split_max = scores.amax(dim=-1) + empty = split_max == float("-inf") + p = torch.exp2(scores - torch.where(empty, 0.0, split_max)[..., None]) + split_sum = p.sum(dim=-1) + p_e4m3 = (p * _P_SCALE).to(torch.float8_e4m3fn).float() + values = keys[..., :KV_LORA_RANK].reshape( + num_tokens, _SPLITS, _SLOTS_PER_SPLIT, KV_LORA_RANK + ) + with true_fp32_matmul(): + split_out = torch.einsum("thrj,trjd->thrd", p_e4m3, values) + merged_max = split_max.amax(dim=-1, keepdim=True) + alpha = torch.where(empty, 0.0, torch.exp2(split_max - merged_max)) + merged_sum = (alpha * split_sum).sum(dim=-1)[..., None] + merged = (alpha[..., None] * split_out).sum(dim=-2) + attn = torch.where( + merged_sum > 0, + merged / (merged_sum * _P_SCALE) * mla_bmm2_scale.reshape(-1)[0].float(), + 0.0, + ) + attn = attn.to(torch.bfloat16).float() + with true_fp32_matmul(): + out = torch.einsum("thd,hvd->thv", attn, v_b_proj.float()) + return out.reshape(num_tokens, NUM_HEADS * V_HEAD_DIM).to(torch.bfloat16) + + def is_valid( + self, + q: torch.Tensor, + kv_pool: torch.Tensor, + indices: torch.Tensor, + seq_lens_kv: torch.Tensor, + mla_bmm1_scale: torch.Tensor, + mla_bmm2_scale: torch.Tensor, + v_b_proj: torch.Tensor, + output: torch.Tensor, + seq_len_q: int, + top_k: int, + block_table: Optional[torch.Tensor] = None, + tokens_per_block: int = 0, + page_index_scale: int = 0, + layer_offset: int = 0, + gather_before_wait: bool = False, + ) -> None: + """The index and table contents the op reads without checking. + + The op checks every shape, dtype and stride loudly; what it trusts is the contents of + ``indices``, ``seq_lens_kv`` and ``block_table``. ``q``, the scales, ``v_b_proj``, + ``output`` and ``gather_before_wait`` are not read here. + """ + num_tokens = indices.shape[0] + device = indices.device + num_pool_rows = kv_pool.numel() // HEAD_DIM_QK + slots = indices[:, :top_k].long() + token = torch.arange(num_tokens, device=device) + seq = token // seq_len_q + kv_len = seq_lens_kv[seq].long() + num_valid = torch.clamp(kv_len - (seq_len_q - 1 - token % seq_len_q), max=top_k) + unread = torch.arange(slots.shape[1], device=device)[None, :] >= num_valid[:, None] + + # Only the first n_t slots of row t are read: a selection packed elsewhere (a valid + # entry behind them) silently drops tokens. + assert not ((slots >= 0) & unread).any(), ( + "indices holds a selected token past the row's valid count " + "min(kv_len - (seq_len_q - 1 - t), top_k); the op never reads it" + ) + selected = (slots >= 0) & ~unread + if block_table is None: + # Global rows are read as given: one past the pool reads other memory. + assert not (selected & (slots >= num_pool_rows)).any(), ( + f"indices names a pool row past the pool's {num_pool_rows} rows" + ) + return + + # The positions are the query token's own context: a later position reads a row this + # step has not appended, or another request's. + own_position = kv_len - seq_len_q + token % seq_len_q + assert not (selected & (slots > own_position[:, None])).any(), ( + "indices selects a position past the query token's own: the row is not this " + "token's context" + ) + # The layer's rows of a slot lie inside the slot: an offset at or past the page scale + # reads the next slot's pages. + assert 0 <= layer_offset < page_index_scale, ( + f"layer_offset {layer_offset} must be below page_index_scale {page_index_scale}" + ) + # Every block of a sequence's KV must be mapped: a narrower table or a -1 inside the + # sequence masks live tokens instead of failing. + num_seqs = num_tokens // seq_len_q + lens = seq_lens_kv[:num_seqs].long() + blocks = (lens + tokens_per_block - 1) // tokens_per_block + width = block_table.shape[1] + assert bool((blocks <= width).all()), ( + f"block_table has {width} columns; a sequence needs {int(blocks.max())}" + ) + live = torch.arange(width, device=device)[None, :] < blocks[:, None] + table = block_table[:num_seqs].long() + assert not (live & (table < 0)).any(), "block_table has a -1 inside a sequence's KV" + if bool(live.any()): + last_page = int(table[live].max()) * page_index_scale + layer_offset + end_row = (last_page + 1) * tokens_per_block + assert end_row <= num_pool_rows, ( + f"block_table maps rows up to {end_row}, past the pool's {num_pool_rows} rows" + ) + + +sparse_mla_flash_decode = GlmSparseMlaFlashDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/_mnnvl_workspace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/_mnnvl_workspace.py new file mode 100644 index 000000000000..9a0cfe6d39ba --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/_mnnvl_workspace.py @@ -0,0 +1,163 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MNNVL Lamport workspaces of the target's hops and exchanges. + +The hops (``mnnvl_hop``, ``mnnvl_moe_finalize_hop``) run on the MNNVL all-reduce workspace of the +engine's ``AllReduce`` (``get_or_scale_allreduce_mnnvl_workspace``): three Lamport buffers whose +empty value is -0.0, and nine uint32 flags (current buffer, dirty buffer, bytes per buffer, dirty +stages, bytes to clear per stage, arrival counter). The hops read, clear and rotate them as the +built-in MNNVL all-reduce does, so both can follow each other on one workspace. + +The workspace is replaced by a larger one when an all-reduce needs more than it holds, and a CUDA +graph keeps the workspace it launches on alive only when it was fetched during that graph's +capture. So the target fetches it for every hop it issues, eagerly and under capture, and never +keeps it across forwards. + +The exchanges that are not all-reduces (``mnnvl_allgather``, the draft steps' argmax exchanges) +each own a ``LamportExchangeWorkspace``, built once and kept. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, + _make_mnnvl_mcast_buffer, + _mnnvl_device_index, + _mnnvl_workspace_all_succeeded, + get_or_scale_allreduce_mnnvl_workspace, +) +from tensorrt_llm.mapping import Mapping + +# Lamport buffers of an MNNVL all-reduce workspace. +NUM_LAMPORT_BUFFERS = 3 + + +def hop_buffer_bytes(num_tokens: int, hidden_size: int, tp_size: int) -> int: + """Bytes of one Lamport buffer that a hop of ``[num_tokens, hidden_size]`` BF16 rows uses.""" + return num_tokens * hidden_size * tp_size * torch.bfloat16.itemsize + + +def hop_supported(num_tokens: int, hidden_size: int, tp_size: int) -> bool: + """Whether the hops serve ``[num_tokens, hidden_size]`` BF16 rows over ``tp_size`` ranks. + + Asked of the current device: 2, 4, 8 or 16 ranks; one-shot (up to 1 MiB of Lamport data per + rank, 1.5 MiB on 16 ranks); one row per cluster of 2 to 8 CTAs of at most 128 threads, eight + elements per thread, covering the row exactly. Hidden 6144 on 16 ranks: 1 to 8 rows. Another + shape runs on the built-in ``AllReduce``. + """ + return bool(torch.ops.trtllm.glm_mnnvl_hop_supported(num_tokens, hidden_size, tp_size)) + + +def _all_reduce_workspace(mapping: Mapping, num_tokens: int, hidden_size: int) -> dict: + return get_or_scale_allreduce_mnnvl_workspace( + mapping, + torch.bfloat16, + buffer_size_bytes=hop_buffer_bytes(num_tokens, hidden_size, mapping.tp_size), + ) + + +def hop_workspace( + mapping: Mapping, num_tokens: int, hidden_size: int +) -> Tuple[torch.Tensor, torch.Tensor]: + """The workspace operands of a hop of ``[num_tokens, hidden_size]`` BF16 rows. + + Returns ``(comm_buffer, buffer_flags)``: the ``[3, N]`` BF16 view of this rank's three Lamport + buffers of the MNNVL all-reduce workspace of ``mapping``, and the workspace's nine uint32 + flags. The workspace grows first when its buffers hold less than the hop needs; it must not + grow while a CUDA graph is captured, so every shape runs once eagerly before its capture. + """ + workspace = _all_reduce_workspace(mapping, num_tokens, hidden_size) + comm_buffer = workspace["uc_buffer"].view(torch.bfloat16).view(NUM_LAMPORT_BUFFERS, -1) + return comm_buffer, workspace["buffer_flags"] + + +def publish_target( + mapping: Mapping, num_tokens: int, hidden_size: int +) -> Tuple[torch.Tensor, torch.Tensor]: + """Where the GEMM that produces a hop's input stores this rank's rows itself. + + Returns ``(lamport, flags)``: the multicast BF16 view (1-D) of the three Lamport buffers of the + workspace ``hop_workspace`` returns for the same shape, and its flags viewed as int32. The GEMM + (the o_proj Lamport publish) reads the current buffer index and the bytes per buffer from the + flags after its grid dependency wait, releases its dependents right after that wait, and stores + its BF16 output, -0.0 written as +0.0, in whole 16-byte vectors at ``[token][tp_rank][column]`` + of the current buffer. The hop that follows it on the stream then runs with + ``shard_published=True`` on that workspace. + """ + workspace = _all_reduce_workspace(mapping, num_tokens, hidden_size) + elements = NUM_LAMPORT_BUFFERS * workspace["buffer_size_bytes"] // torch.bfloat16.itemsize + lamport = workspace["handle"].get_mc_buffer((elements,), torch.bfloat16, 0) + return lamport, workspace["buffer_flags"].view(torch.int32) + + +class LamportExchangeWorkspace: + """A private MNNVL Lamport workspace of the TP group, for an exchange that is not an all-reduce. + + Three Lamport buffers whose unwritten 32-bit words hold 0xFFFFFFFF, so that every payload, + -0.0 included, travels unchanged, and their nine flags in the layout of the all-reduce + workspace's. Used by ``mnnvl_allgather`` (the verify logits) and by the draft-step kernels + that exchange each rank's argmax (index, value) packet, with the token-table row and the + hidden-half share where the next eh_proj comes from the token table (``draft_sample_glue``, + ``draft_ehproj_table``). Each of those kernels empties the buffer of its previous call with + its own threads, so every user owns its workspace, sized for its largest call. + + Construction is collective over the TP group: every rank constructs the same workspaces in the + same order, outside CUDA graph capture (the target does it once, in ``post_load_weights``), + and a failure on any rank raises on every rank. + + Args: + mapping: The tensor-parallel mapping of the group. + buffer_size_bytes: Bytes of one Lamport buffer, a positive multiple of 16. + + Attributes: + comm_buffer: This rank's unicast view of the three buffers as 1-D int32, the kernels' + ``comm_buffer`` (``exchange_buffer``) operand. + buffer_flags: The nine uint32 Lamport flags, the kernels' ``buffer_flags`` + (``exchange_flags``) operand. + """ + + def __init__(self, mapping: Mapping, buffer_size_bytes: int): + if buffer_size_bytes <= 0 or buffer_size_bytes % 16: + raise ValueError( + f"buffer_size_bytes must be a positive multiple of 16: {buffer_size_bytes}" + ) + comm = _get_mnnvl_workspace_comm(mapping) + workspace_bytes = NUM_LAMPORT_BUFFERS * buffer_size_bytes + state = None + error: Optional[Exception] = None + try: + handle = _make_mnnvl_mcast_buffer( + comm, workspace_bytes, mapping, mapping.is_multi_node() + ) + comm_buffer = handle.get_uc_buffer( + mapping.tp_rank, (workspace_bytes // 4,), torch.int32, 0 + ) + # Filled under inference mode, like the all-reduce workspace: the engine may build + # it while CUDA graph warm-up runs in inference mode. + with torch.inference_mode(): + comm_buffer.fill_(-1) + buffer_flags = torch.tensor( + [0, 2, buffer_size_bytes, 0, 0, 0, 0, 0, 0], + dtype=torch.uint32, + device=torch.device("cuda", _mnnvl_device_index(mapping)), + ) + torch.cuda.synchronize() + state = (handle, comm_buffer, buffer_flags) + except Exception as e: + error = e + # Also fences the fill: no rank stores into a peer's buffer before the peer filled it. + if not _mnnvl_workspace_all_succeeded(comm, error is None): + raise RuntimeError( + "MNNVL Lamport exchange workspace construction failed on at least one rank" + ) from error + assert state is not None + # The handle keeps the multicast allocation mapped; the communicator stays with it. + self._handle, self.comm_buffer, self.buffer_flags = state + self._comm = comm + self.mapping = mapping + self.buffer_size_bytes = buffer_size_bytes diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_allgather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_allgather.py new file mode 100644 index 000000000000..bd2144dd6fea --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_allgather.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MNNVL all-gather of BF16 rows widened to FP32, via the trtllm glm_mnnvl_allgather op.""" + +from __future__ import annotations + +from typing import Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._mnnvl_workspace import NUM_LAMPORT_BUFFERS + + +class GlmMnnvlAllgather(OpWrapper): + """Every rank's BF16 rows gathered along the last dimension in rank order, as FP32. + + ``out[..., r * shard:(r + 1) * shard]`` is rank ``r``'s ``input`` widened to FP32, the same on + every rank: the all-gather, concatenation and ``.float()`` of the vocabulary-parallel LM head's + logits, in one kernel on a private ``_mnnvl_workspace.LamportExchangeWorkspace``. Each rank + stores its shard once through the multicast address and copies the gathered rows out of its + local buffer as they land. Bit-exact, except that a pair of BF16 values whose 32 bits are all + set (NaNs no conversion produces) arrives with its lowest bit cleared, still NaNs. + + The kernel waits for its producer before it reads ``input`` and releases its PDL dependents + right after that wait: a dependent reads the output only after its own grid dependency wait. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The c1 verify step at TP16: 8 rows of the LM head's vocabulary shard. + CELLS: tuple[Cell, ...] = ( + Cell( + why="verify-step logits: 8 rows of this rank's 9,680-entry vocabulary shard " + "(154,880 over 16 ranks)", + spec=dict(rows=8, shard=9680), + ), + ) + + note = """ + Every rank calls the gather in the same order with the same shape. The gathered rows of one + call must fit one Lamport buffer of the workspace (8 rows of the 154,880-entry vocabulary are + 2.5 MB of BF16); the op refuses more. The workspace is the gather's alone: the kernel empties + the buffer of its previous call with its own threads. + """ + + def __call__( + self, + input: torch.Tensor, + comm_buffer: torch.Tensor, + buffer_flags: torch.Tensor, + group_size: int, + ) -> torch.Tensor: + return torch.ops.trtllm.glm_mnnvl_allgather(input, comm_buffer, buffer_flags, group_size) + + def reference( + self, + per_rank_inputs: Sequence[torch.Tensor], + comm_buffer: torch.Tensor | None = None, + buffer_flags: torch.Tensor | None = None, + group_size: int | None = None, + ) -> torch.Tensor: + """The rank-ordered concatenation along the last dimension, widened to FP32. + + A collective: ``per_rank_inputs`` holds every rank's ``input`` in rank order, in place of + this rank's. The other arguments are ignored. + """ + return torch.cat([rows.float() for rows in per_rank_inputs], dim=-1) + + def is_valid( + self, + input: torch.Tensor, + comm_buffer: torch.Tensor, + buffer_flags: torch.Tensor, + group_size: int, + ) -> None: + # The all-reduce workspace marks empty words with -0.0 instead of 0xFFFFFFFF: on it the + # gather takes stale words for rows that have not arrived. + assert comm_buffer.dtype == torch.int32 and comm_buffer.dim() == 1, ( + "comm_buffer must be the comm_buffer of a LamportExchangeWorkspace" + ) + # Flags of another workspace make the kernel poll and fill buffers whose state they do + # not describe: stale words are gathered, or a buffer still in use is emptied. + bytes_per_buffer = int(buffer_flags.view(torch.int32)[2].item()) + workspace_bytes = comm_buffer.numel() * comm_buffer.element_size() + assert NUM_LAMPORT_BUFFERS * bytes_per_buffer == workspace_bytes, ( + "buffer_flags must be the flags of comm_buffer's workspace" + ) + + +mnnvl_allgather = GlmMnnvlAllgather() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_hop.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_hop.py new file mode 100644 index 000000000000..a9a591882939 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_hop.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One-shot MNNVL all-reduce with the residual add and RMSNorm, via the trtllm glm_mnnvl_hop op.""" + +from __future__ import annotations + +from typing import Sequence, Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._mnnvl_workspace import NUM_LAMPORT_BUFFERS + + +class GlmMnnvlHop(OpWrapper): + """All-reduce of BF16 rows over the ranks plus ``residual``, and the RMSNorm of that sum. + + ``residual_out = sum of input over the ranks + residual``, ``norm_out = RMSNorm(residual_out)``, + in one kernel on the MNNVL all-reduce workspace (``_mnnvl_workspace.hop_workspace``): every rank + stores its rows through the multicast address into its slot of the current Lamport buffer of + every rank and polls all slots in one loop. The sum is in FP32 in rank order (from zero per + chunk of 8 ranks, the chunks added in order) and rounded to BF16; the residual is added in + BF16 and the RMSNorm normalizes that BF16 sum, ``(x * rsqrt(mean(x^2) + epsilon)) * gamma``. + Every rank gets the same bits. Returns ``(norm_out, residual_out)``, both fresh. + + The kernel releases its PDL dependents right after its grid dependency wait (at its start when + the shard is published): a dependent reads ``norm_out`` and ``residual_out`` only after its + own grid dependency wait. + + With ``shard_published`` the GEMM that produced ``input`` (o_proj, through + ``_mnnvl_workspace.publish_target``) has already stored this rank's rows into the workspace; + ``input`` gives the shape only and is not read. That GEMM must be the hop's predecessor on the + stream and release its dependents only after its own grid dependency wait. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The c1 decode step at TP16: 8 rows on a verify step and draft step 0, 1 row on draft steps + # 1-6; every rank holds the full hidden size 6144. The world size is the test's. + CELLS: tuple[Cell, ...] = ( + Cell( + why="post-attention hop of the 78 verify layers and draft step 0: 8 rows, the shard " + "published by o_proj", + spec=dict(rows=8, hidden=6144, shard_published=True, zero_residual=False), + ), + Cell( + why="post-MLP hop of the dense layers 0-2: 8 rows, the shard in input", + spec=dict(rows=8, hidden=6144, shard_published=False, zero_residual=False), + ), + Cell( + why="draft step 0's eh_proj hop: input_layernorm with a zero residual, the sum as " + "residual_out", + spec=dict(rows=8, hidden=6144, shard_published=False, zero_residual=True), + ), + Cell( + why="post-attention hop of draft steps 1-6: one row, the shard published by o_proj", + spec=dict(rows=1, hidden=6144, shard_published=True, zero_residual=False), + ), + ) + + note = """ + Every rank of the group calls the hop in the same order, with the same shape and the same + `shard_published`; a disagreement hangs the group. `_mnnvl_workspace.hop_supported` states the + shapes the op takes (hidden 6144 on 16 ranks: 1 to 8 rows); it refuses others, which the target + runs on the built-in `AllReduce`. A -0.0 among the summed elements arrives as +0.0, because + -0.0 marks an empty Lamport slot. + """ + + def __call__( + self, + input: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + epsilon: float, + comm_buffer: torch.Tensor, + buffer_flags: torch.Tensor, + shard_published: bool, + ) -> Tuple[torch.Tensor, torch.Tensor]: + norm_out, residual_out = torch.ops.trtllm.glm_mnnvl_hop( + input, residual, gamma, epsilon, comm_buffer, buffer_flags, shard_published + ) + return norm_out, residual_out + + def reference( + self, + per_rank_inputs: Sequence[torch.Tensor], + residual: torch.Tensor, + gamma: torch.Tensor, + epsilon: float, + comm_buffer: torch.Tensor | None = None, + buffer_flags: torch.Tensor | None = None, + shard_published: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(norm_out, residual_out)`` in FP32, each rounded once. + + A collective: ``per_rank_inputs`` holds every rank's ``input`` in rank order, in place of + this rank's. ``residual`` is the calling rank's. ``comm_buffer``, ``buffer_flags`` and + ``shard_published`` decide how the rows travel, not what is computed, and are ignored. + """ + total = torch.zeros_like(residual, dtype=torch.float32) + for rows in per_rank_inputs: + total = total + rows.float() + summed = total + residual.float() + rms = torch.rsqrt(summed.pow(2).mean(-1, keepdim=True) + epsilon) + norm = summed * rms * gamma.float() + return norm.to(residual.dtype), summed.to(residual.dtype) + + def is_valid( + self, + input: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + epsilon: float, + comm_buffer: torch.Tensor, + buffer_flags: torch.Tensor, + shard_published: bool, + ) -> None: + # The kernel reads the slots as BF16 and takes -0.0 as the empty value: a workspace with + # another empty value (the 0xFFFFFFFF words of an exchange workspace) hands it stale words + # as rows. + assert comm_buffer.dtype == torch.bfloat16, ( + "comm_buffer must be the BF16 view of the MNNVL all-reduce workspace (hop_workspace)" + ) + # Flags of another workspace make the kernel poll, clear and rotate buffers whose state + # they do not describe: stale slots are summed, or a buffer still in use is emptied. + bytes_per_buffer = int(buffer_flags.view(torch.int32)[2].item()) + workspace_bytes = comm_buffer.numel() * comm_buffer.element_size() + assert NUM_LAMPORT_BUFFERS * bytes_per_buffer == workspace_bytes, ( + "buffer_flags must be the flags of comm_buffer's workspace" + ) + + +mnnvl_hop = GlmMnnvlHop() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_moe_finalize_hop.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_moe_finalize_hop.py new file mode 100644 index 000000000000..00bba97f581e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_moe_finalize_hop.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MoE finalize fused into the post-MoE MNNVL hop, via trtllm glm_mnnvl_moe_finalize_hop.""" + +from __future__ import annotations + +from typing import Sequence, Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._mnnvl_workspace import NUM_LAMPORT_BUFFERS + + +class GlmMnnvlMoeFinalizeHop(OpWrapper): + """``mnnvl_hop`` of the finalized MoE rows: the routed sum plus ``addend``, all-reduced. + + On each rank, row ``t`` of the finalized rows is ``addend[t] + sum over k of + expert_weights[t, k] * gemm2_output[expanded_idx_to_permuted_idx[t, k]]``, over the slots whose + index is not -1: the routed sum in FP32 with an FMA in ``k`` order, rounded to BF16, then the + addend added in FP32 and rounded once more, as the TRTLLM-Gen finalize with the shared + experts' output as addend computes it. The rows are then all-reduced with the residual add and + RMSNorm exactly as ``mnnvl_hop`` does (without a published shard), on the same workspace. + Returns ``(norm_out, residual_out)``, both fresh, the same bits on every rank. + + The kernel reads ``expert_weights`` and ``expanded_idx_to_permuted_idx`` before its grid + dependency wait, so the kernel that wrote them must have completed when this one launches: the + MoE's routing kernel, when this hop follows FC2 on the stream and FC2 releases its dependents + only after its own grid dependency wait; or a constant written outside graph capture. The + other operands are read after the wait. The PDL dependents launch right after that wait and + read the outputs only after their own. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The c1 decode step at TP16, hidden 6144, every expert local (moe_ep 1). + CELLS: tuple[Cell, ...] = ( + Cell( + why="post-MoE hop of the 75 MoE layers of a verify step: 8 rows, top-8 FC2 rows, the " + "shared experts' output as addend", + spec=dict(rows=8, hidden=6144, routing="top8"), + ), + Cell( + why="post-MoE hop of an MTP draft step: one row, its routed row in slot 0 with " + "weight 1 and slots 1-7 empty (-1), the shared experts' output as addend", + spec=dict(rows=1, hidden=6144, routing="routed_row"), + ), + ) + + note = """ + The shapes are those of `mnnvl_hop` (`_mnnvl_workspace.hop_supported`), with exactly 8 routing + slots per token; the same group discipline applies. FC2 rows may be wider than the hidden size + (`gemm2_output.shape[1]` is the row pitch); the first `hidden` columns are read. + """ + + def __call__( + self, + gemm2_output: torch.Tensor, + expert_weights: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + addend: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + epsilon: float, + comm_buffer: torch.Tensor, + buffer_flags: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + norm_out, residual_out = torch.ops.trtllm.glm_mnnvl_moe_finalize_hop( + gemm2_output, + expert_weights, + expanded_idx_to_permuted_idx, + addend, + residual, + gamma, + epsilon, + comm_buffer, + buffer_flags, + ) + return norm_out, residual_out + + def reference( + self, + per_rank_inputs: Sequence[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]], + residual: torch.Tensor, + gamma: torch.Tensor, + epsilon: float, + comm_buffer: torch.Tensor | None = None, + buffer_flags: torch.Tensor | None = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(norm_out, residual_out)`` in FP32, each rounded once. + + A collective: ``per_rank_inputs`` holds every rank's ``(gemm2_output, expert_weights, + expanded_idx_to_permuted_idx, addend)`` in rank order, in place of this rank's. + ``residual`` is the calling rank's. ``comm_buffer`` and ``buffer_flags`` are ignored. + """ + hidden = residual.shape[-1] + total = torch.zeros_like(residual, dtype=torch.float32) + for gemm2_output, expert_weights, permuted_idx, addend in per_rank_inputs: + taken = permuted_idx >= 0 + rows = gemm2_output[permuted_idx.clamp(min=0).long(), :hidden].float() + weights = torch.where(taken, expert_weights.float(), 0.0) + total = total + (weights.unsqueeze(-1) * rows).sum(dim=1) + addend.float() + summed = total + residual.float() + rms = torch.rsqrt(summed.pow(2).mean(-1, keepdim=True) + epsilon) + norm = summed * rms * gamma.float() + return norm.to(residual.dtype), summed.to(residual.dtype) + + def is_valid( + self, + gemm2_output: torch.Tensor, + expert_weights: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + addend: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + epsilon: float, + comm_buffer: torch.Tensor, + buffer_flags: torch.Tensor, + ) -> None: + # An index outside [-1, rows) reads memory past the FC2 rows as an expert row. + index = expanded_idx_to_permuted_idx + in_range = (index >= -1) & (index < gemm2_output.shape[0]) + assert bool(in_range.all()), ( + "expanded_idx_to_permuted_idx must index gemm2_output's rows, or be -1" + ) + # As for mnnvl_hop: a workspace with another empty value, or flags of another workspace, + # make the kernel sum stale slots or empty a buffer still in use. + assert comm_buffer.dtype == torch.bfloat16, ( + "comm_buffer must be the BF16 view of the MNNVL all-reduce workspace (hop_workspace)" + ) + bytes_per_buffer = int(buffer_flags.view(torch.int32)[2].item()) + workspace_bytes = comm_buffer.numel() * comm_buffer.element_size() + assert NUM_LAMPORT_BUFFERS * bytes_per_buffer == workspace_bytes, ( + "buffer_flags must be the flags of comm_buffer's workspace" + ) + + +mnnvl_moe_finalize_hop = GlmMnnvlMoeFinalizeHop() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits.py new file mode 100644 index 000000000000..71df8ceeac9e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Compile cache and launch of the indexer's FP8 paged MQA logits kernel (FP32 epilogue). + +The kernel (``_fp8_paged_mqa_logits_kernel.FP8MQALogitsKernel``) scores every query row of a +request against the request's paged indexer keys: ``logits[b * next_n + t, j] = scale[j] * +sum_h relu(q[b, t, h] . k[j]) * weights[b * next_n + t, h]`` for the keys ``j`` of the request's +context. Two policies decide the compiled variant: + +- full query: a request's ``next_n`` query rows stay one atom of ``next_n * num_heads`` MMA + columns (``full_query_weight_cache_limit``); the weight register cache is capped so that the + shared-memory tail handles the remaining heads without changing the FMA order; +- early KV (``early_kv``, shapes of ``early_kv_supported``): the N-split variant launched with + PDL, which loads the older KV blocks before ``griddepcontrol.wait``. + +All variants compute the same per-element arithmetic in the same order. +""" + +from __future__ import annotations + +import functools +from typing import Dict, Optional, Tuple + +import cutlass +import cutlass.cute as cute +import torch + +from ._fp8_paged_mqa_logits_kernel import FP8MQALogitsKernel + +#: Keys per MMA tile; a schedule split is one tile per math warp group. +COMPUTE_BLOCK_KV = 128 +#: Keys per schedule split (DeepGEMM's SPLIT_KV of the SM100 paged metadata). +SPLIT_KV = 2 * COMPUTE_BLOCK_KV + +_KERNELS: Dict[Tuple, object] = {} + + +@functools.lru_cache(maxsize=None) +def num_sms(device_index: int) -> int: + """SMs of the device, the number of persistent CTAs the kernel launches.""" + return torch.cuda.get_device_properties(device_index).multi_processor_count + + +def full_query_weight_cache_limit( + next_n: int, num_heads: int, head_dim: int, phys_block_kv: int, q_dtype: torch.dtype +) -> Optional[int]: + """The weight register cache of the full-query kernel, or None for the default kernel. + + The two MMA groups share 512 TMEM columns. Each needs ``next_n * num_heads`` columns, with an + eight-column MMA granularity. D = 128 and three Q/KV stages fit the shared-memory budget + throughout this range. Physical pages must divide the 128-token compute tile into at most four + pieces. At most 96 FP32 weights per thread (16 for each query row) stay in registers; the + shared-memory tail handles the remaining heads without changing the FMA order. Up to four + query rows keep the default kernel. + """ + columns = next_n * num_heads + if ( + next_n <= 4 + or head_dim != 128 + or not 8 <= columns <= 256 + or columns % 8 + or phys_block_kv not in (32, 64, 128) + or num_heads % 4 + or q_dtype != torch.float8_e4m3fn + ): + return None + limit = min(16, num_heads, (96 // next_n // 4) * 4) + return limit if limit >= 4 else None + + +def early_kv_supported(next_n: int, num_heads: int) -> bool: + """Shapes of the early-KV kernel (PDL launch, N split): each group's accumulator splits into + two N halves of next_n / 2 query rows, and a half must be a whole number of 16-column UMMA + atoms.""" + return next_n % 2 == 0 and (next_n * num_heads // 2) % 16 == 0 + + +def _compile( + phys_block_kv: int, + num_heads: int, + head_dim: int, + next_n: int, + sms: int, + weight_cache_limit: Optional[int], + early_kv: bool, +): + """Compile one variant against fake tensors with symbolic runtime extents (TVM FFI).""" + n = next_n * num_heads + block_bytes = phys_block_kv * (head_dim + 4) + sym_num_phys_blocks = cute.sym_int() + sym_b = cute.sym_int() + max_ctx = cute.sym_int() + max_blocks_per_seq = cute.sym_int() + num_ctas = cute.sym_int() + + # The indexer K-cache pool view can be strided in dim 0 (pages of other layers in between): + # the per-page stride is read at runtime; the bytes of one page are contiguous. + kv_fake = cute.runtime.make_fake_tensor( + cutlass.Uint8, (sym_num_phys_blocks, block_bytes), stride=(cute.sym_int64(), 1) + ) + q_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Uint8, (n, head_dim, sym_b), stride_order=(1, 0, 2) + ) + w_fake = cute.runtime.make_fake_compact_tensor(cutlass.Float32, (n, sym_b), stride_order=(0, 1)) + logits_fake = cute.runtime.make_fake_tensor( + cutlass.Float32, (cute.sym_int(), max_ctx), stride=(cute.sym_int64(), 1) + ) + bt_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_b, max_blocks_per_seq), stride_order=(1, 0) + ) + cl_fake = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (sym_b,), stride_order=(0,)) + sm_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (num_ctas, 2), stride_order=(1, 0) + ) + fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + + kernel = FP8MQALogitsKernel( + block_kv=COMPUTE_BLOCK_KV, + phys_block_kv=phys_block_kv, + num_heads=num_heads, + head_dim=head_dim, + next_n=next_n, + num_sms=sms, + num_epi_subtiles=1, + epi_dtype=cutlass.Float32, + acc_dtype=cutlass.Float32, + output_dtype=cutlass.Float32, + weight_cache_limit=weight_cache_limit, + split_n=early_kv, + early_kv=early_kv, + ) + return cute.compile( + kernel, + kv_fake, + q_fake, + w_fake, + logits_fake, + bt_fake, + cl_fake, + sm_fake, + cutlass.Int32(1), + cutlass.Int32(1), + fake_stream, + options="--enable-tvm-ffi", + ) + + +def run_fp8_paged_mqa_logits( + q: torch.Tensor, + kv_cache: torch.Tensor, + weights: torch.Tensor, + context_lens: torch.Tensor, + block_table: torch.Tensor, + schedule_meta: torch.Tensor, + logits: torch.Tensor, + max_context_len: int, + early_kv: bool, +) -> torch.Tensor: + """Launch the kernel; returns ``logits[:, :max_context_len]``. + + Args: + q: [B, next_n, H, D] float8_e4m3fn. + kv_cache: [num_pages, page_size, 1, D + 4] uint8 indexer K-cache view. + weights: [B * next_n, H] float32. + context_lens: [B] int32 keys per request. + block_table: [B, max_pages] int32, contiguous. + schedule_meta: [num_sms + 1, 2] int32 (the SM100 paged MQA-logits schedule). + logits: [B * next_n, >= round_up(max_context_len, SPLIT_KV)] float32 output, unit + column stride; the kernel writes whole splits. + max_context_len: columns of the returned view. + early_kv: run the early-KV variant where its shape allows (``early_kv_supported``). + """ + b, next_n, h, d = q.shape + n = next_n * h + phys_block_kv = kv_cache.shape[1] + num_phys_blocks = kv_cache.shape[0] + sms = num_sms(q.device.index if q.device.index is not None else torch.cuda.current_device()) + weight_cache_limit = full_query_weight_cache_limit(next_n, h, d, phys_block_kv, q.dtype) + early_kv = early_kv and early_kv_supported(next_n, h) + key = (phys_block_kv, h, d, next_n, sms, weight_cache_limit, early_kv) + compiled = _KERNELS.get(key) + if compiled is None: + compiled = _KERNELS[key] = _compile( + phys_block_kv, h, d, next_n, sms, weight_cache_limit, early_kv + ) + # Q [B, next_n, H, D] -> [N, D, B]; weights [B * next_n, H] -> [N, B]; FP8 Q as uint8 bytes. + q_3d = q.reshape(b, n, d).permute(1, 2, 0).view(torch.uint8) + w_2d = weights.reshape(b, n).t() + kv_flat = kv_cache.reshape(num_phys_blocks, -1) + out = logits[:, :max_context_len] + compiled(kv_flat, q_3d, w_2d, out, block_table, context_lens, schedule_meta, num_phys_blocks, b) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits_kernel.py new file mode 100644 index 000000000000..ca300c385d6e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits_kernel.py @@ -0,0 +1,2088 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +""" +CuTe DSL FP8 paged MQA logits kernel for Blackwell (SM100 family), with the N split and the early KV +loads of the GLM indexer decode. + +Architecture: + - 384 threads: 256 math (2 WGs) + 128 specialized (2 TMA + 2 UMMA) + - 1 TMA per KV block [128, 128], UMMA iterates 4x K=32 + - 2 warp groups process 2 KV blocks per iteration (kNumMathWarpGroups=2) + - Q reloaded via TMA pipeline when q_idx (batch) changes + - Persistent kernel: CTAs iterate through assigned (q_idx, kv_idx) pairs + - Weights cached in registers: preloaded once per q_idx change (not per KV block) + - KV Scales loaded via TMA to SMEM (separate pipeline per group, Math consumes) + +Merged KV+Scale pipeline: + - KV data and scales share a single TMA barrier per group + - TMA loads both KV and Scale under one barrier (combined tx_count) + - UMMA waits on merged barrier (for KV GEMM), does NOT release + - Math waits on merged barrier (for scale read), Math releases + +Fused KV layout: + - KV data and scales stored contiguously per physical block: + [num_phys_blocks, block_kv * (head_dim + 4)] bytes + - Per block: [KV_all_tokens (block_kv * head_dim bytes)] [Scales (block_kv * 4 bytes)] + - KV and Scale views are derived inside __call__ using CuTE pointer arithmetic + +Scheduler: + - schedule_meta[sm_idx] = (start_q_idx, start_kv_idx / kNumMathWarpGroups) + - schedule_meta[sm_idx+1] = end boundary for this CTA + - fetch_next_task pattern: each warp role independently advances (q_idx, kv_idx) + - kv_idx in units of KV blocks, advances by kNumMathWarpGroups=2 per step + +N split (split_n): + - A group's accumulator of one KV block (next_n * num_heads columns) is computed in two + N halves (query rows [0, next_n/2) and [next_n/2, next_n)), each its own TMEM stage + with its own UMMA commit + - Q is loaded as one TMA box per half; the UMMA warp issues the block's K slices half by half + - The math warpgroup drains and releases each half on its own (double-buffered TMEM loads), + so the MMAs of the next block's first half overlap the epilogue of the current block's + second half + - Per-element epilogue math and the FMA order of every row are unchanged + +Early KV loads (early_kv): + - Launched with programmatic dependent launch + - Before griddepcontrol.wait the TMA warps read the schedule, context lengths and block + table and load up to num_kv_stages KV blocks per group that end before the request's + newest next_n tokens + - Q, weights and every later KV block are loaded after the wait; nothing is written before it + - The caller guarantees that the grids this launch depends on write no KV row older than + the newest next_n of each request, and no block table, context length or schedule entry + +Dynamic shape support: + - Model-constant dims (block_kv, head_dim, N, per_token) remain static for codegen + - Runtime-varying dims (batch_size, num_phys_blocks, max_ctx, max_blocks_per_seq, + num_ctas) are marked dynamic via mark_compact_shape_dynamic + - Allows JIT cache reuse across different batch sizes / sequence lengths + +Epilogue dtype flows (--acc_dtype / --epi_dtype): + + Flow 1: --acc_dtype fp32 --epi_dtype fp16 + Q(FP8) x K(FP8) -> MMA acc(FP32) -> TMEM(FP32) + -> LDTM -> Reg(FP32) -> cvt FP16 -> ReLU(FP16) + -> FMA(fma.rn.f16x2) with weights(FP16 from SMEM) -> partial sum(FP16) + -> x scale(FP32->FP16) -> cvt output_dtype -> store logits(output_dtype) + + Flow 2: --acc_dtype fp16 --epi_dtype fp16 + Q(FP8) x K(FP8) -> MMA acc(FP16) -> TMEM(FP16, pack_16b) + -> LDTM -> Reg(FP16) -> ReLU(FP16) + -> FMA(fma.rn.f16x2) with weights(FP16 from SMEM) -> partial sum(FP16) + -> x scale(FP32->FP16) -> cvt output_dtype -> store logits(output_dtype) + + --output_dtype: fp32 (default), fp16, bf16. Controls logits tensor dtype and final store conversion. + Default: --acc_dtype fp32 --epi_dtype fp32 --output_dtype fp32 +""" + +from typing import Tuple + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass import Float16, Int32 +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm, vector +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +# CuTe DSL CUDA 13 validates rounding modes as string literals. The string +# form is also accepted by older wrappers, so keep it version-independent. +_RND_RN = "rn" + + +@dsl_user_op +def pack_f16x2( + a: Float16, + b: Float16, + *, + loc=None, + ip=None, +) -> Int32: + f16_ty = Float16.mlir_type + i32_ty = Int32.mlir_type + vec2_f16 = ir.VectorType.get([2], f16_ty, loc=loc) + v = vector.from_elements( + vec2_f16, + (Float16(a).ir_value(loc=loc, ip=ip), Float16(b).ir_value(loc=loc, ip=ip)), + loc=loc, + ip=ip, + ) + return Int32(llvm.bitcast(i32_ty, v, loc=loc, ip=ip)) + + +@dsl_user_op +def unpack_f16x2( + packed: Int32, + *, + loc=None, + ip=None, +) -> Tuple[Float16, Float16]: + f16_ty = Float16.mlir_type + vec2_f16 = ir.VectorType.get([2], f16_ty, loc=loc) + v = llvm.bitcast(vec2_f16, Int32(packed).ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + r0 = Float16(vector.extract(v, dynamic_position=[], static_position=[0], loc=loc, ip=ip)) + r1 = Float16(vector.extract(v, dynamic_position=[], static_position=[1], loc=loc, ip=ip)) + return r0, r1 + + +@dsl_user_op +def fma_f16x2( + a: Int32, + b: Int32, + c: Int32, + *, + loc=None, + ip=None, +) -> Int32: + i32_ty = Int32.mlir_type + return Int32( + llvm.inline_asm( + i32_ty, + [ + Int32(a).ir_value(loc=loc, ip=ip), + Int32(b).ir_value(loc=loc, ip=ip), + Int32(c).ir_value(loc=loc, ip=ip), + ], + "fma.rn.f16x2 $0, $1, $2, $3;", + "=r,r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def max_f16x2( + a: Int32, + b: Int32, + *, + loc=None, + ip=None, +) -> Int32: + i32_ty = Int32.mlir_type + return Int32( + llvm.inline_asm( + i32_ty, + [Int32(a).ir_value(loc=loc, ip=ip), Int32(b).ir_value(loc=loc, ip=ip)], + "max.f16x2 $0, $1, $2;", + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def add_f16x2( + a: Int32, + b: Int32, + *, + loc=None, + ip=None, +) -> Int32: + i32_ty = Int32.mlir_type + return Int32( + llvm.inline_asm( + i32_ty, + [Int32(a).ir_value(loc=loc, ip=ip), Int32(b).ir_value(loc=loc, ip=ip)], + "add.f16x2 $0, $1, $2;", + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +class FP8MQALogitsKernel: + """FP8 paged MQA logits kernel for Blackwell (SM100). + + Each CTA processes a range of (q_idx, kv_split) pairs. + A split = 2 consecutive KV blocks within a sequence (one per warp group). + Q is shared between warp groups and reloaded when q_idx changes. + """ + + def __init__( + self, + block_kv: int = 128, + phys_block_kv: int = 128, + num_heads: int = 64, + head_dim: int = 128, + next_n: int = 1, + num_sms: int = 148, + remove_kv_wait_in_epilogue: bool = False, + early_tmem_copy: bool = False, + smem_subpartition_opt: bool = False, + max_kv_pipeline: bool = False, + max_umma_pipeline: bool = False, + num_epi_subtiles: int = 1, + epi_dtype=cutlass.Float32, + acc_dtype=cutlass.Float32, + output_dtype=cutlass.Float32, + weight_cache_limit: int | None = None, + split_n: bool = False, + early_kv: bool = False, + ): + self.block_kv = block_kv + self.phys_block_kv = phys_block_kv + self.num_blocks_per_mma = block_kv // phys_block_kv + assert block_kv % phys_block_kv == 0, ( + f"block_kv={block_kv} must be divisible by phys_block_kv={phys_block_kv}" + ) + assert self.num_blocks_per_mma <= 4, ( + f"num_blocks_per_mma={self.num_blocks_per_mma} exceeds max 4" + ) + self.remove_kv_wait_in_epilogue = remove_kv_wait_in_epilogue + self.early_tmem_copy = early_tmem_copy + self.smem_subpartition_opt = smem_subpartition_opt + self.num_heads = num_heads + if weight_cache_limit is not None and (weight_cache_limit <= 0 or weight_cache_limit % 4): + raise ValueError("weight_cache_limit must be a positive multiple of four") + self.weight_cache_limit = num_heads if weight_cache_limit is None else weight_cache_limit + self.head_dim = head_dim + self.next_n = next_n + self.N = next_n * num_heads + self.num_sms = num_sms + self.num_epi_subtiles = num_epi_subtiles + self.epi_dtype = epi_dtype + self.epi_bytes = 2 if epi_dtype == cutlass.Float16 else 4 + # sW stage stride padded to 128-byte SMEM alignment for TMA bulk copy. + # Without padding, e.g. fp16 + N=32 gives 64B per stage, so stage 1 + # at +64 would be misaligned (TMA requires 128-byte aligned SMEM dest). + w_stage_bytes = self.N * self.epi_bytes + self.w_stage_stride = ((w_stage_bytes + 127) // 128 * 128) // self.epi_bytes + self.output_dtype = output_dtype + if num_epi_subtiles > 1 and num_heads % num_epi_subtiles != 0: + raise ValueError("num_heads must be divisible by num_epi_subtiles") + if (num_heads // num_epi_subtiles) % 4 != 0: + raise ValueError( + "num_heads // num_epi_subtiles must be divisible by 4 (FMA unroll granularity)" + ) + if split_n and ( + next_n % 2 != 0 + or (self.N // 2) % 16 != 0 + or num_epi_subtiles != 1 + or epi_dtype != cutlass.Float32 + or acc_dtype != cutlass.Float32 + ): + raise ValueError( + "split_n needs an even next_n, N / 2 divisible by 16, one epilogue sub-tile " + "and an FP32 accumulator and epilogue" + ) + self.split_n = split_n + self.n_halves = 2 if split_n else 1 + self.early_kv = early_kv + self.num_groups = 2 + + self.num_math_threads = 256 + self.num_specialized_threads = 128 + self.threads_per_cta = 384 + self.num_math_warps = 8 + self.tma_warp_base = 8 + self.umma_warp_base = 10 + + self.num_q_stages = 3 # 3 stages for Q pipelining across batch sequences + + # TMEM: 512 columns total, each group needs N columns per UMMA stage + # max_umma_stages = 512 // (2 * N) + TMEM_COLS = 512 + if split_n: + # One TMEM stage per N half: 2 groups x 2 halves x N / 2 columns. + self.num_umma_stages = self.n_halves + elif max_umma_pipeline: + self.num_umma_stages = min(2, TMEM_COLS // (2 * self.N)) + else: + self.num_umma_stages = 1 + + if max_kv_pipeline: + smem_capacity = utils.get_smem_capacity_in_bytes() + # Reserve ~1 KB for barriers and misc + SMEM_BUDGET = smem_capacity - 1024 + # KV+Scale per stage (×2 groups): + # 2 * (block_kv * head_dim * 1B + block_kv * 4B) + kv_scale_per_stage = 2 * (block_kv * head_dim + block_kv * 4) + # Q+W per stage: Q is N * head_dim * 1B, W uses padded stride + qw_per_stage = self.N * head_dim + self.w_stage_stride * self.epi_bytes + qw_total = qw_per_stage * self.num_q_stages + self.num_kv_stages = (SMEM_BUDGET - qw_total) // kv_scale_per_stage + else: + self.num_kv_stages = 3 + + # Pad SMEM to push sW/sScales into sub-partition 1 (>= 128KB), + # avoiding sub-bank conflicts with UMMA reading sKV. + # Layout: barriers(~256B) | sKV_0 | sKV_1 | sQ | [pad] | sW | sScales + if self.smem_subpartition_opt: + BOUNDARY = 128 * 1024 + used = ( + 256 + + 2 * (block_kv * head_dim * self.num_kv_stages) + + self.N * head_dim * self.num_q_stages + ) + used = ((used + 127) // 128) * 128 + if used < BOUNDARY: + self.smem_pad_bytes = ((BOUNDARY - used + 1023) // 1024) * 1024 + else: + self.smem_pad_bytes = 0 + else: + self.smem_pad_bytes = 0 + + self.acc_dtype = acc_dtype + self.cta_group = tcgen05.CtaGroup.ONE + self.cluster_shape_mn = (1, 1) + self.mma_tiler_mn = (block_kv, self.N // self.n_halves) + + def _setup_mma(self, a_dtype, b_dtype, a_major, b_major): + self.a_dtype = a_dtype + self.b_dtype = b_dtype + self.a_major_mode = a_major + self.b_major_mode = b_major + + self.mma_tiler = (*self.mma_tiler_mn, 1) + tiled_mma = sm100_utils.make_trivial_tiled_mma( + a_dtype, + a_major, + b_major, + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) # 32 + + # Full-K: tile K = head_dim (128), 1 TMA per block + mma_inst_tile_k = self.head_dim // mma_inst_shape_k # 4 + full_k = mma_inst_shape_k * mma_inst_tile_k # 128 + self.mma_tiler = ( + self.mma_tiler_mn[0], + self.mma_tiler_mn[1], + full_k, + ) + + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.epi_tile = self.cta_tile_shape_mnk[:2] + + # KV SMEM: 3 stages per group, each stage holds full [128, 128] + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + a_dtype, + self.num_kv_stages, + ) + # Q SMEM: num_q_stages stages of [N, 128]; with split_n, Q stage s holds its two + # N halves as B stages 2s and 2s + 1. + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + b_dtype, + self.num_q_stages * self.n_halves, + ) + + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape) + self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake) + self.num_tmem_alloc_cols_total = ( + self.num_tmem_alloc_cols * self.num_groups * self.num_umma_stages + ) + + return tiled_mma + + def _q_b_stage(self, q_stage, hh): + """B (Q) SMEM stage of N half hh of Q stage q_stage.""" + if self.n_halves == 1: + return q_stage + return q_stage * self.n_halves + hh + + def _copy_q(self, q_copy, q_bar, q_stage, q_idx): + """TMA copies of request q_idx's Q (one box per N half) and weights into Q stage q_stage.""" + tma_atom_b, tBgB_h, tBsB, b_mcast_mask, tma_atom_w, tWgW, tWsW = q_copy + for hh in range(self.n_halves): + cute.copy( + tma_atom_b, + tBgB_h[hh][(None, 0, q_idx)], + tBsB[(None, self._q_b_stage(q_stage, hh))], + tma_bar_ptr=q_bar, + mcast_mask=b_mcast_mask, + ) + cute.copy( + tma_atom_w, + tWgW[(None, q_idx)], + tWsW[(None, q_stage)], + tma_bar_ptr=q_bar, + ) + + def _row_sum(self, acc_vec, t, w_cache, sW, q_stage, num_w_in_reg): + """sum_h w[t, h] * relu(acc[h]) of query row t, in the FP32 epilogue's FMA order.""" + zero = cutlass.Float32(0.0) + s0x, s0y, s1x, s1y = zero, zero, zero, zero + for h in range(0, self.num_heads, 4): + a0 = cutlass.max(acc_vec[h], zero) + a1 = cutlass.max(acc_vec[h + 1], zero) + a2 = cutlass.max(acc_vec[h + 2], zero) + a3 = cutlass.max(acc_vec[h + 3], zero) + if h < num_w_in_reg: + r0 = t * num_w_in_reg + h + w0, w1, w2, w3 = w_cache[r0], w_cache[r0 + 1], w_cache[r0 + 2], w_cache[r0 + 3] + else: + c0 = t * self.num_heads + h + w0 = sW[(c0, q_stage)] + w1 = sW[(c0 + 1, q_stage)] + w2 = sW[(c0 + 2, q_stage)] + w3 = sW[(c0 + 3, q_stage)] + s0x, s0y = cute.arch.fma_packed_f32x2((a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN) + s1x, s1y = cute.arch.fma_packed_f32x2((a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN) + return s0x + s0y + s1x + s1y + + @cute.jit + def __call__( + self, + kv_fused: cute.Tensor, # Fused KV: [num_phys_blocks, block_bytes] FP8 + b: cute.Tensor, # Q: [N, head_dim, batch_size] + weights: cute.Tensor, # [N, batch_size] (transposed for TMA) + logits: cute.Tensor, # [batch_size * next_n, max_context_len] + block_table: cute.Tensor, # [batch_size, max_blocks_per_seq] + context_lens: cute.Tensor, # [batch_size] + schedule_meta: cute.Tensor, # [num_sms+1, 2] int32 + num_phys_blocks: cutlass.Int32, + batch_size: cutlass.Int32, + stream: cuda.CUstream, + ): + # Derive KV and Scale views from fused buffer using CuTE ops. + # Fused layout per physical block: [KV data (phys_block_kv*head_dim)] [Scales (phys_block_kv*4)] + phys_block_kv = self.phys_block_kv + scale_offset_elems = phys_block_kv * self.head_dim # in FP8 elements + + # Recast fused buffer to FP8 (same 1-byte elements, needed for MMA type inference) + kv_fp8 = cute.recast_tensor(kv_fused, cutlass.Float8E4M3FN) + + # Q (b) was passed as uint8 to work around DLPack's lack of float8 support; + # recast back to FP8 so MMA type inference and TMA descriptors are correct. + b = cute.recast_tensor(b, cutlass.Float8E4M3FN) + + # Read the real per-block stride (bytes = FP8 elements) from the input. + # When KV is the indexer K-cache pool view, the pool is laid out as + # [num_blocks, num_layers, kvFactor, blockSize], so dim-0 stride = + # num_layers * kvFactor * phys_block_bytes (not phys_block_bytes). + # Using the input stride keeps both contiguous test path and the + # strided prod path correct. FP8 = 1 byte per element, so the byte + # stride is the same as the element stride after recast. + kv_block_stride = kv_fused.layout.stride[0] + + # KV view: [phys_block_kv, head_dim, num_phys_blocks] FP8 + # Each TMA loads one physical block; multiple TMAs fill a compute tile. + kv_layout = cute.make_layout( + (phys_block_kv, self.head_dim, num_phys_blocks), + stride=(self.head_dim, 1, kv_block_stride), + ) + a = cute.make_tensor(kv_fp8.iterator, kv_layout) + + # Scale view: offset pointer to scale region, recast FP8 → Float32 + # [phys_block_kv, num_phys_blocks] float32 (after recast) + scale_fp8_layout = cute.make_layout( + (phys_block_kv * 4, num_phys_blocks), + stride=(1, kv_block_stride), + ) + scale_fp8 = cute.make_tensor(kv_fp8.iterator + scale_offset_elems, scale_fp8_layout) + scales = cute.recast_tensor(scale_fp8, cutlass.Float32) + + a_dtype = a.element_type + b_dtype = b.element_type + a_major = utils.LayoutEnum.from_tensor(a).mma_major_mode() + b_major = utils.LayoutEnum.ROW_MAJOR.mma_major_mode() + + tiled_mma = self._setup_mma(a_dtype, b_dtype, a_major, b_major) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # TMA for KV (A) — fmha_decode_paged pattern. + # Build a TMA SMEM layout via tiled_divide on the full compute-tile + # layout, then select to drop trivial K dim. Atom uses mode [0] as + # single-tile SMEM layout and (phys, head) as cta_tiler. + tma_load_op = cpasync.CopyBulkTensorTileG2SOp() + self.a_tma_view_layout = sm100_utils.make_smem_layout( + tcgen05.OperandMajorMode.K, + (self.block_kv, self.head_dim), + a_dtype, + self.num_kv_stages, + ) + self.a_tma_view_layout = cute.tiled_divide( + self.a_tma_view_layout, (self.phys_block_kv, self.head_dim) + ) + # ((tile_M, tile_K), rest_M, rest_K, stages) → drop trivial rest_K + self.a_tma_view_layout = cute.select(self.a_tma_view_layout, mode=[0, 1, 3]) + # ((tile_M, tile_K), rest_M=num_sub_blocks, stages) + tma_atom_a, tma_tensor_a = cpasync.make_tiled_tma_atom( + tma_load_op, + a, + self.a_tma_view_layout[0], # atom SMEM = single-tile (mode 0) + (self.phys_block_kv, self.head_dim), + ) + + # TMA for Q (B) — full K=128, L dim = batch_size (unchanged) + b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA for Weights — [N, batch_size], tile [N], L=batch_size + self.w_smem_layout_staged = cute.make_layout( + (self.N, self.num_q_stages), + stride=(1, self.w_stage_stride), + ) + w_smem_per_stage = cute.select(self.w_smem_layout_staged, mode=[0]) + tma_atom_w, tma_tensor_w = cpasync.make_tiled_tma_atom( + tma_load_op, + weights, + w_smem_per_stage, + self.w_smem_layout_staged.shape[:1], + ) + + # TMA for Scales — [phys_block_kv, num_phys_blocks], tile [phys_block_kv] + # SMEM holds compute_block_kv scales per stage; filled by + # num_blocks_per_mma sub-block TMAs at consecutive offsets. + self.s_smem_layout_staged = cute.make_layout((self.block_kv, self.num_kv_stages)) + s_smem_per_subblock = cute.make_layout((phys_block_kv,)) + tma_atom_s, tma_tensor_s = cpasync.make_tiled_tma_atom( + tma_load_op, + scales, + s_smem_per_subblock, + (phys_block_kv,), + ) + + b_copy_size = cute.size_in_bytes(b_dtype, b_smem_layout) + w_copy_size = self.N * self.epi_bytes + # Per sub-block: phys_block_kv * head_dim (KV) + phys_block_kv * 4 (scales) + kv_tma_bytes_per_subblock = phys_block_kv * self.head_dim + scale_tma_bytes_per_subblock = phys_block_kv * 4 + # Total per compute tile = num_blocks_per_mma sub-blocks + self.num_kv_scale_tma_bytes = self.num_blocks_per_mma * ( + kv_tma_bytes_per_subblock + scale_tma_bytes_per_subblock + ) + # Q + Weights share barrier (like DeepGEMM); one Q box per N half + self.num_q_tma_bytes = b_copy_size * atom_thr_size * self.n_halves + w_copy_size + + num_ctas = self.num_sms + + @cute.struct + class SharedStorage: + kv_mbar_0: cute.struct.MemRange[cutlass.Int64, self.num_kv_stages * 2] + kv_mbar_1: cute.struct.MemRange[cutlass.Int64, self.num_kv_stages * 2] + q_mbar: cute.struct.MemRange[cutlass.Int64, self.num_q_stages * 2] + umma_mbar_0: cute.struct.MemRange[cutlass.Int64, self.num_umma_stages * 2] + umma_mbar_1: cute.struct.MemRange[cutlass.Int64, self.num_umma_stages * 2] + tmem_holding_buf: cutlass.Int32 + + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_w, + tma_tensor_w, + tma_atom_s, + tma_tensor_s, + logits, + block_table, + context_lens, + schedule_meta, + batch_size, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.w_smem_layout_staged, + self.s_smem_layout_staged, + self.a_tma_view_layout, + self.epi_tile, + SharedStorage, + ).launch( + grid=(1, 1, num_ctas), + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + use_pdl=self.early_kv, + ) + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, # KV pool + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, # Q (L dim = batch_size) + tma_atom_w: cute.CopyAtom, + mW_tma: cute.Tensor, # Weights TMA coord tensor [N, batch_size] + tma_atom_s: cute.CopyAtom, + mS_tma: cute.Tensor, # Scales TMA coord tensor [block_kv, num_phys_blocks] + mLogits: cute.Tensor, # [batch_size * next_n, max_context_len] + mBlockTable: cute.Tensor, # [batch_size, max_blocks_per_seq] + mContextLens: cute.Tensor, # [batch_size] + mScheduleMeta: cute.Tensor, # [num_sms+1, 2] int32 + batch_size: cutlass.Int32, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + w_smem_layout_staged: cute.Layout, + s_smem_layout_staged: cute.Layout, + a_tma_view_layout: cute.ComposedLayout, + epi_tile: cute.Tile, + SharedStorage: cutlass.Constexpr, + ): + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + tidx, _, _ = cute.arch.thread_idx() + + # Warp roles (matches DeepGEMM SM100) + warpgroup_idx = warp_idx // 4 + is_math_warp = warp_idx < 8 + is_tma_warp_0 = warp_idx == 8 + is_tma_warp_1 = warp_idx == 9 + is_tma_warp = is_tma_warp_0 | is_tma_warp_1 + is_umma_warp_0 = warp_idx == 10 + is_umma_warp_1 = warp_idx == 11 + + # Early schedule metadata load: issue global loads ASAP so their + # ~200-cycle L2 latency overlaps with subsequent prologue setup + # (SMEM alloc, TMA partition, MMA fragment creation, etc.) + NUM_MATH_WG = 2 # kNumMathWarpGroups + NUM_BLOCKS_PER_MMA = self.num_blocks_per_mma + sm_idx = bidz + start_q = mScheduleMeta[(sm_idx, 0)] + start_kv_half = mScheduleMeta[(sm_idx, 1)] + end_q_idx = mScheduleMeta[(sm_idx + 1, 0)] + end_kv_half = mScheduleMeta[(sm_idx + 1, 1)] + # Early mContextLens load: overlap ~200-cycle L2 latency with the + # entire prologue setup (pipelines, SMEM alloc, TMA partition, etc.) + # Clamp to avoid OOB when start_q == batch_size (zero-work CTA sentinel). + # Note: zero-work CTAs get a stale current_num_kv (from the last batch + # element), but it is never used because has_work will be False. + start_q_clamped = min(start_q, batch_size - 1) + current_num_kv = (mContextLens[start_q_clamped] + self.block_kv - 1) // self.block_kv + if cutlass.const_expr(self.early_kv): + current_ctx_tok = mContextLens[start_q_clamped] + + if is_tma_warp: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + cpasync.prefetch_descriptor(tma_atom_w) + cpasync.prefetch_descriptor(tma_atom_s) + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + block_kv_val = self.block_kv + num_heads = self.num_heads + next_n = self.next_n + num_epi_subtiles = self.num_epi_subtiles + + # === Pipelines === + prod_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + + # Q pipeline: TMA producer → Math consumer (8 math warps) + # PipelineTmaAsync: consumer_release uses is_signalling_thread + # (lane 0 per warp). 8 math warps × 1 lane-0 = 8 arrives. + q_cons_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 8) + q_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.q_mbar.data_ptr(), + num_stages=self.num_q_stages, + producer_group=prod_group, + consumer_group=q_cons_group, + tx_count=self.num_q_tma_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=tidx, + defer_sync=True, + ) + q_prod_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_q_stages + ) + # Both Math WGs share the same pipeline state (advance in lockstep) + q_cons_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_q_stages + ) + # UMMA warps observe Q pipeline (wait only, no release) + # to ensure Q is in SMEM before GEMM. Critical for UMMA warp 1 + # since TMA warp 1 only loads KV1 (not Q). + q_cons_state_umma_0 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_q_stages + ) + q_cons_state_umma_1 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_q_stages + ) + + # Merged KV+Scale pipelines (per-group, 3 stages each) + # Like DeepGEMM: KV data and scales share one barrier. + # TMA loads both under one barrier. Math is consumer (releases). + # UMMA also waits on this barrier (for KV GEMM) but does NOT release. + math_warps_per_group = self.num_math_warps // 2 # 4 warps + kv_cons_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, math_warps_per_group) + kv_pipeline_0 = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.kv_mbar_0.data_ptr(), + num_stages=self.num_kv_stages, + producer_group=prod_group, + consumer_group=kv_cons_group, + tx_count=self.num_kv_scale_tma_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=tidx, + defer_sync=True, + ) + kv_pipeline_1 = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.kv_mbar_1.data_ptr(), + num_stages=self.num_kv_stages, + producer_group=prod_group, + consumer_group=kv_cons_group, + tx_count=self.num_kv_scale_tma_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=tidx, + defer_sync=True, + ) + + kv_prod_state_0 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_kv_stages + ) + kv_prod_state_1 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_kv_stages + ) + # UMMA consumer states (wait only, no release) + kv_cons_state_umma_0 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_kv_stages + ) + kv_cons_state_umma_1 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_kv_stages + ) + # Math consumer states (wait + release) + kv_cons_state_math_0 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_kv_stages + ) + kv_cons_state_math_1 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_kv_stages + ) + + # UMMA pipelines (per-group) + math_threads_per_group = self.num_math_threads // 2 + umma_pipeline_0 = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.umma_mbar_0.data_ptr(), + num_stages=self.num_umma_stages, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, math_threads_per_group), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + umma_pipeline_1 = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.umma_mbar_1.data_ptr(), + num_stages=self.num_umma_stages, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, math_threads_per_group), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + umma_prod_state_0 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_umma_stages + ) + umma_prod_state_1 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_umma_stages + ) + umma_cons_state_0 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_umma_stages + ) + umma_cons_state_1 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_umma_stages + ) + + # TMEM — only Math warps (8×32=256) + UMMA warps (2×32=64) = 320 threads + # TMA warps do NOT participate, so they can start TMA loads earlier. + # Math warp 0 is the allocator (like fp16_gemm_3's epilogue warp 0), + # because math warps are the last TMEM consumers (epilogue reads). + tmem_alloc_num_threads = 320 # 10 warps: warp 0-7 (math) + warp 10-11 (umma) + tmem_alloc_barrier = pipeline.NamedBarrier(barrier_id=1, num_threads=tmem_alloc_num_threads) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=0, # math warp 0 does alloc+free (last TMEM consumer) + is_two_cta=False, + ) + + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # SMEM allocation: per-group KV + shared Q + sKV_0 = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + sKV_1 = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + sQ = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + # Pad SMEM to push sW/sScales into sub-partition 1 (>= 128KB) + # to avoid sub-bank conflicts with UMMA reading sKV from + # sub-partition 0. + if cutlass.const_expr(self.smem_pad_bytes > 0): + _ = smem.allocate(self.smem_pad_bytes) + # Weights SMEM: [N, num_q_stages], shared Q barrier + sW = smem.allocate_tensor( + element_type=self.epi_dtype, + layout=w_smem_layout_staged, + byte_alignment=128, + ) + # Scales SMEM: [block_kv, num_kv_stages] float32, per group + sScales_0 = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=s_smem_layout_staged, + byte_alignment=128, + ) + sScales_1 = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=s_smem_layout_staged, + byte_alignment=128, + ) + + a_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + # Partition KV (A): fmha_decode_paged pattern. + # SMEM view is ((tile), num_sub_blocks, stages) — built in __call__. + # Use .outer (plain layout); swizzle is captured by sKV_0's iterator. + # GMEM: local_tile by (phys, head), then group first 2 modes into tile. + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + sKV_0_for_tma = cute.make_tensor(sKV_0.iterator, a_tma_view_layout.outer) + sKV_1_for_tma = cute.make_tensor(sKV_1.iterator, a_tma_view_layout.outer) + gA = cute.local_tile( + mA_mkl, + (self.phys_block_kv, self.head_dim), + coord=(None, None, None), + ) + tAsA_0, tAgA_0 = cpasync.tma_partition( + tma_atom_a, + 0, + cute.make_layout(1), + sKV_0_for_tma, + cute.group_modes(gA, 0, 2), + ) + tAsA_1, tAgA_1 = cpasync.tma_partition( + tma_atom_a, + 0, + cute.make_layout(1), + sKV_1_for_tma, + cute.group_modes(gA, 0, 2), + ) + + # Partition Q (B): shared SMEM, L dim = batch_size + gB_nkl = cute.local_tile( + mB_nkl, + cute.slice_(self.mma_tiler, (0, None, None)), + (None, None, None), + ) + tCgB = thr_mma.partition_B(gB_nkl) + b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sQ, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + # [tma, K, L] per N half + tBgB_h = [tBgB[(None, hh, None, None)] for hh in range(self.n_halves)] + + # Partition Weights: standalone TMA, [N, batch_size] → [N] per stage + w_cta_layout = cute.make_layout((1,)) + tWsW, tWgW = cpasync.tma_partition( + tma_atom_w, + 0, + w_cta_layout, + cute.group_modes(sW, 0, 1), + cute.group_modes(mW_tma, 0, 1), + ) + q_copy = (tma_atom_b, tBgB_h, tBsB, b_mcast_mask, tma_atom_w, tWgW, tWsW) + + # Partition Scales: explicit sub_blocks + stages dims. + # Layout (phys_block_kv, num_sub, stages) K-major with custom strides. + s_tma_view_layout = cute.make_layout( + (self.phys_block_kv, self.num_blocks_per_mma, self.num_kv_stages), + stride=(1, self.phys_block_kv, self.block_kv), + ) + sScales_0_for_tma = cute.make_tensor(sScales_0.iterator, s_tma_view_layout) + sScales_1_for_tma = cute.make_tensor(sScales_1.iterator, s_tma_view_layout) + # GMEM: local_tile by phys to match atom's tile size + gS = cute.local_tile(mS_tma, (self.phys_block_kv,), coord=(None, None)) + tSsS_0, tSgS_0 = cpasync.tma_partition( + tma_atom_s, + 0, + cute.make_layout(1), + sScales_0_for_tma, + gS, + ) + tSsS_1, tSgS_1 = cpasync.tma_partition( + tma_atom_s, + 0, + cute.make_layout(1), + sScales_1_for_tma, + gS, + ) + + # MMA fragments + tCrA_0 = tiled_mma.make_fragment_A(sKV_0) + tCrA_1 = tiled_mma.make_fragment_A(sKV_1) + tCrB = tiled_mma.make_fragment_B(sQ) # shared + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + + # Staged acc (fp16_gemm_3 pattern): append UMMA stage dim + # shape: (*acc_shape, STAGE) — dynamic index on last dim reduces rank + us = self.num_umma_stages + cols = self.num_tmem_alloc_cols + acc_shape_staged = cute.append(acc_shape, us) + tCtAcc_fake_staged = tiled_mma.make_fragment_C(acc_shape_staged) + + # TMEM layout info (allocation deferred to UMMA/Math warp branches) + cols_per_group = cols * us * (32 // self.acc_dtype.width) + num_tmem_alloc_cols_total = self.num_tmem_alloc_cols_total + + # Epilogue setup + c_layout = utils.LayoutEnum.ROW_MAJOR + epi_sub_mn = (epi_tile[0], num_heads // num_epi_subtiles) + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + c_layout, + self.acc_dtype, + self.acc_dtype, + epi_sub_mn, + use_2cta_instrs, + ) + + # ===== SCHEDULER: derive values from early-loaded schedule metadata ===== + end_kv_idx = end_kv_half * NUM_MATH_WG + + # Convert start to KV block units (matching DeepGEMM) + current_q_idx = start_q + current_kv_idx = start_kv_half * NUM_MATH_WG + + # ===== COMMON SCHEDULER STATE (before warp branches) ===== + # Each warp role independently maintains its own copy of these + # variables (like DeepGEMM where each role creates its own scheduler). + # Pre-fetch first task (current_num_kv loaded early above for latency hiding) + next_q_idx = current_q_idx + next_kv_idx = current_kv_idx + next_num_kv = current_num_kv + if cutlass.const_expr(self.early_kv): + # TMA warps: context length of next_q_idx, and the CTA's first request + next_ctx_tok = current_ctx_tok + first_q = current_q_idx + # Sentinel: no previous batch (matches DeepGEMM's q_idx = batch_size) + q_idx = batch_size + # While-loop termination flag (matches DeepGEMM's fetch_next_task pattern). + # True if this CTA has work assigned (start != end in schedule_meta). + has_work = (current_q_idx != end_q_idx) | (current_kv_idx != end_kv_idx) + + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # ===== WARP-SPECIALIZED EXECUTION ===== + + if is_tma_warp_0: + # TMA warp 0: loads Q (prefetch) + KV for group 0 + # Matches DeepGEMM's TMA warp with kv_group_idx == 0 + cute.arch.warpgroup_reg_dealloc(24) + lane_idx = tidx % 32 + + # Block table prefetch: 32 lanes cache block indices, + # distributed via shuffle. Each lane holds num_blocks_per_mma + # physical block indices per compute tile. (DeepGEMM L233-244) + cached_blks = [cutlass.Int32(0) for _ in range(NUM_BLOCKS_PER_MMA)] + kv_blk_ptr = cutlass.Int32(32) # force prefetch on first use + + if cutlass.const_expr(not self.early_kv): + # Prefetch first Q before loop (like DeepGEMM line 203-204) + q_pipeline.producer_acquire(q_prod_state) + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + self._copy_q(q_copy, q_bar, q_prod_state.index, next_q_idx) + q_prod_state.advance() + else: + # Q is loaded after griddepcontrol.wait. Until then, count the KV + # tiles loaded; the first iteration's Q prefetch waits for it too. + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + waited = cutlass.Int32(0) + pre_wait_tiles = cutlass.Int32(0) + ran = cutlass.Int32(0) + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + if cutlass.const_expr(self.early_kv): + if waited == 0: + # Wait before the first tile that holds one of the request's newest + # next_n tokens, before a KV stage would be reused, and at a request change. + must_wait = cutlass.Int32(0) + if (kv_idx + 1) * block_kv_val > next_ctx_tok - next_n: + must_wait = cutlass.Int32(1) + if pre_wait_tiles >= self.num_kv_stages: + must_wait = cutlass.Int32(1) + if q_idx_old != batch_size: + if q_idx != q_idx_old: + must_wait = cutlass.Int32(1) + if must_wait != 0: + cute.arch.griddepcontrol_wait() + waited = cutlass.Int32(1) + q_pipeline.producer_acquire(q_prod_state) + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + self._copy_q(q_copy, q_bar, q_prod_state.index, first_q) + q_prod_state.advance() + if q_idx_old != batch_size: + # The Q prefetch the first iteration deferred + deferred_q = first_q + 1 + if deferred_q < end_q_idx: + q_pipeline.producer_acquire(q_prod_state) + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + self._copy_q(q_copy, q_bar, q_prod_state.index, deferred_q) + q_prod_state.advance() + elif deferred_q == end_q_idx: + if end_kv_idx > 0: + q_pipeline.producer_acquire(q_prod_state) + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + self._copy_q(q_copy, q_bar, q_prod_state.index, deferred_q) + q_prod_state.advance() + + # Q prefetch: when batch changes, load Q for NEXT batch + if q_idx != q_idx_old: + kv_blk_ptr = cutlass.Int32(32) # force re-prefetch + prefetch_next = q_idx + 1 + if cutlass.const_expr(self.early_kv): + if waited == 0: + # Deferred until the wait + prefetch_next = end_q_idx + 1 + if prefetch_next < end_q_idx: + q_pipeline.producer_acquire(q_prod_state) + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + self._copy_q(q_copy, q_bar, q_prod_state.index, prefetch_next) + q_prod_state.advance() + elif prefetch_next == end_q_idx: + if end_kv_idx > 0: + q_pipeline.producer_acquire(q_prod_state) + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + self._copy_q(q_copy, q_bar, q_prod_state.index, prefetch_next) + q_prod_state.advance() + + # Block table prefetch for group 0 (like DeepGEMM L233-241). + # Each lane loads num_blocks_per_mma physical block indices + # for one compute tile (kv_idx counts compute tiles). + if kv_blk_ptr == 32: + kv_blk_ptr = cutlass.Int32(0) + prefetch_kv = kv_idx + lane_idx * NUM_MATH_WG + if prefetch_kv < num_kv: + base_phys = prefetch_kv * NUM_BLOCKS_PER_MMA + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + cached_blks[i] = mBlockTable[(q_idx, base_phys + i)] + else: + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + cached_blks[i] = cutlass.Int32(0) + + # Get block indices via shuffle before barrier (like DeepGEMM L244) + phys_blks = [cutlass.Int32(0)] * NUM_BLOCKS_PER_MMA + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + phys_blks[i] = cute.arch.shuffle_sync(cached_blks[i], kv_blk_ptr) + kv_blk_ptr = kv_blk_ptr + 1 + + # Load KV + Scale for group 0: num_blocks_per_mma TMAs per tile. + kv_pipeline_0.producer_acquire(kv_prod_state_0) + bar = kv_pipeline_0.producer_get_barrier(kv_prod_state_0) + stage = kv_prod_state_0.index + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + cute.copy( + tma_atom_a, + tAgA_0[(None, 0, 0, phys_blks[i])], + tAsA_0[(None, i, stage)], + tma_bar_ptr=bar, + mcast_mask=a_mcast_mask, + ) + cute.copy( + tma_atom_s, + tSgS_0[(None, 0, phys_blks[i])], + tSsS_0[(None, i, stage)], + tma_bar_ptr=bar, + ) + kv_prod_state_0.advance() + if cutlass.const_expr(self.early_kv): + ran = cutlass.Int32(1) + if waited == 0: + pre_wait_tiles = pre_wait_tiles + 1 + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = (mContextLens[next_q_idx] + block_kv_val - 1) // block_kv_val + if cutlass.const_expr(self.early_kv): + next_ctx_tok = mContextLens[next_q_idx] + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + if cutlass.const_expr(self.early_kv): + if waited == 0: + # Every tile of this CTA was loaded before the wait (or it has none). + cute.arch.griddepcontrol_wait() + q_pipeline.producer_acquire(q_prod_state) + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + self._copy_q(q_copy, q_bar, q_prod_state.index, first_q) + q_prod_state.advance() + if ran != 0: + # The Q prefetch the first iteration deferred + deferred_q = first_q + 1 + if deferred_q < end_q_idx: + q_pipeline.producer_acquire(q_prod_state) + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + self._copy_q(q_copy, q_bar, q_prod_state.index, deferred_q) + q_prod_state.advance() + elif deferred_q == end_q_idx: + if end_kv_idx > 0: + q_pipeline.producer_acquire(q_prod_state) + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + self._copy_q(q_copy, q_bar, q_prod_state.index, deferred_q) + q_prod_state.advance() + + elif is_tma_warp_1: + # TMA warp 1: loads KV + Scale for group 1 only + # Matches DeepGEMM's TMA warp with kv_group_idx == 1 + cute.arch.warpgroup_reg_dealloc(24) + lane_idx = tidx % 32 + + # Block table prefetch for group 1 + cached_blks = [cutlass.Int32(0) for _ in range(NUM_BLOCKS_PER_MMA)] + kv_blk_ptr = cutlass.Int32(32) # force prefetch on first use + if cutlass.const_expr(self.early_kv): + waited = cutlass.Int32(0) + pre_wait_tiles = cutlass.Int32(0) + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + if cutlass.const_expr(self.early_kv): + if waited == 0: + # Same rule as TMA warp 0, for tile kv_idx + 1 + must_wait = cutlass.Int32(0) + if (kv_idx + 2) * block_kv_val > next_ctx_tok - next_n: + must_wait = cutlass.Int32(1) + if pre_wait_tiles >= self.num_kv_stages: + must_wait = cutlass.Int32(1) + if q_idx_old != batch_size: + if q_idx != q_idx_old: + must_wait = cutlass.Int32(1) + if must_wait != 0: + cute.arch.griddepcontrol_wait() + waited = cutlass.Int32(1) + + # New q_idx → force block table re-prefetch + if q_idx != q_idx_old: + kv_blk_ptr = cutlass.Int32(32) + + # Block table prefetch for group 1 (like DeepGEMM L233-241). + if kv_blk_ptr == 32: + kv_blk_ptr = cutlass.Int32(0) + prefetch_kv = kv_idx + 1 + lane_idx * NUM_MATH_WG + if prefetch_kv < num_kv: + base_phys = prefetch_kv * NUM_BLOCKS_PER_MMA + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + cached_blks[i] = mBlockTable[(q_idx, base_phys + i)] + else: + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + cached_blks[i] = cutlass.Int32(0) + + # Get block indices via shuffle before barrier (like DeepGEMM L244) + phys_blks = [cutlass.Int32(0)] * NUM_BLOCKS_PER_MMA + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + phys_blks[i] = cute.arch.shuffle_sync(cached_blks[i], kv_blk_ptr) + kv_blk_ptr = kv_blk_ptr + 1 + + # Load KV + Scale for group 1: num_blocks_per_mma TMAs per tile. + kv_pipeline_1.producer_acquire(kv_prod_state_1) + bar = kv_pipeline_1.producer_get_barrier(kv_prod_state_1) + stage = kv_prod_state_1.index + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + cute.copy( + tma_atom_a, + tAgA_1[(None, 0, 0, phys_blks[i])], + tAsA_1[(None, i, stage)], + tma_bar_ptr=bar, + mcast_mask=a_mcast_mask, + ) + cute.copy( + tma_atom_s, + tSgS_1[(None, 0, phys_blks[i])], + tSsS_1[(None, i, stage)], + tma_bar_ptr=bar, + ) + kv_prod_state_1.advance() + if cutlass.const_expr(self.early_kv): + if waited == 0: + pre_wait_tiles = pre_wait_tiles + 1 + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = (mContextLens[next_q_idx] + block_kv_val - 1) // block_kv_val + if cutlass.const_expr(self.early_kv): + next_ctx_tok = mContextLens[next_q_idx] + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + elif is_umma_warp_0: + # UMMA warp for group 0 + # Must wait on Q pipeline: TMA operations with different + # barriers are NOT visibility-ordered even within the same + # warp. KV0 barrier arriving does not guarantee Q SMEM + # writes are visible. + cute.arch.warpgroup_reg_dealloc(24) + + # TMEM: wait for math warp 0's allocation, retrieve pointer + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base_0 = cute.make_tensor(tmem_ptr, tCtAcc_fake_staged.layout) + tCtAcc_base_1 = cute.make_tensor(tmem_ptr + cols_per_group, tCtAcc_fake_staged.layout) + + if is_leader_cta: + num_k_blocks = cute.size(tCrA_0.shape[2]) + q_stage_0 = cutlass.Int32(0) + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + # Wait for Q pipeline when batch changes + if q_idx != q_idx_old: + if q_idx_old < batch_size: + q_cons_state_umma_0.advance() + q_pipeline.consumer_wait(q_cons_state_umma_0) + q_stage_0 = q_cons_state_umma_0.index + + # Process KV block for group 0 (kv_idx + 0) + # Unconditional UMMA (like DeepGEMM): OOB iterations + # compute on garbage data; results written to aligned + # padding region in logits buffer. + # Wait KV first, then TMEM empty (like DeepGEMM) + kv_pipeline_0.consumer_wait(kv_cons_state_umma_0) + # One accumulator stage per N half (a single one without split_n) + for hh in cutlass.range_constexpr(self.n_halves): + umma_pipeline_0.producer_acquire(umma_prod_state_0) + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + kv_stage = kv_cons_state_umma_0.index + q_b_stage = self._q_b_stage(q_stage_0, hh) + tCtAcc_0 = tCtAcc_base_0[(None, None, None, umma_prod_state_0.index)] + for k_block in cutlass.range_constexpr(num_k_blocks): + cute.gemm( + tiled_mma, + tCtAcc_0, + tCrA_0[None, None, k_block, kv_stage], + tCrB[None, None, k_block, q_b_stage], + tCtAcc_0, + ) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + if cutlass.const_expr(hh == self.n_halves - 1): + # No consumer_release here — Math WG 0 releases + kv_cons_state_umma_0.advance() + + umma_pipeline_0.producer_commit(umma_prod_state_0) + umma_prod_state_0.advance() + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = ( + mContextLens[next_q_idx] + block_kv_val - 1 + ) // block_kv_val + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + elif is_umma_warp_1: + # UMMA warp for group 1 + # Explicitly waits on Q pipeline — critical because TMA warp 1 + # only loads KV1, not Q. Without this wait, UMMA warp 1 can + # start GEMM before TMA warp 0 finishes loading Q into SMEM. + cute.arch.warpgroup_reg_dealloc(24) + + # TMEM: wait for umma_warp_0's allocation, retrieve pointer + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base_0 = cute.make_tensor(tmem_ptr, tCtAcc_fake_staged.layout) + tCtAcc_base_1 = cute.make_tensor(tmem_ptr + cols_per_group, tCtAcc_fake_staged.layout) + + if is_leader_cta: + num_k_blocks_1 = cute.size(tCrA_1.shape[2]) + q_stage_1 = cutlass.Int32(0) + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + # Wait for Q pipeline when batch changes + if q_idx != q_idx_old: + if q_idx_old < batch_size: + q_cons_state_umma_1.advance() + q_pipeline.consumer_wait(q_cons_state_umma_1) + q_stage_1 = q_cons_state_umma_1.index + + # Process KV block for group 1 (kv_idx + 1) + # Unconditional UMMA (like DeepGEMM) + # Wait KV first, then TMEM empty (like DeepGEMM) + kv_pipeline_1.consumer_wait(kv_cons_state_umma_1) + # One accumulator stage per N half (a single one without split_n) + for hh in cutlass.range_constexpr(self.n_halves): + umma_pipeline_1.producer_acquire(umma_prod_state_1) + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + kv_stage_1 = kv_cons_state_umma_1.index + q_b_stage = self._q_b_stage(q_stage_1, hh) + tCtAcc_1 = tCtAcc_base_1[(None, None, None, umma_prod_state_1.index)] + for k_block in cutlass.range_constexpr(num_k_blocks_1): + cute.gemm( + tiled_mma, + tCtAcc_1, + tCrA_1[None, None, k_block, kv_stage_1], + tCrB[None, None, k_block, q_b_stage], + tCtAcc_1, + ) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + if cutlass.const_expr(hh == self.n_halves - 1): + # No consumer_release here — Math WG 1 releases + kv_cons_state_umma_1.advance() + + umma_pipeline_1.producer_commit(umma_prod_state_1) + umma_prod_state_1.advance() + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = ( + mContextLens[next_q_idx] + block_kv_val - 1 + ) // block_kv_val + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + elif is_math_warp: + cute.arch.warpgroup_reg_alloc(240) + + # TMEM: math warp 0 is the allocator; all math warps wait + retrieve + tmem.allocate(num_tmem_alloc_cols_total) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base_0 = cute.make_tensor(tmem_ptr, tCtAcc_fake_staged.layout) + tCtAcc_base_1 = cute.make_tensor(tmem_ptr + cols_per_group, tCtAcc_fake_staged.layout) + + local_tidx = tidx % 128 + cC = cute.make_identity_tensor(epi_sub_mn) + + if cutlass.const_expr(self.split_n): + # Per KV block: wait for each N half's accumulator, drain its next_n / 2 rows + # (double-buffered TMEM loads) and release it; the KV stage is released once the + # MMAs of both halves are complete. + rows_per_half = next_n // self.n_halves + # Weight register cache as the FP32 epilogue below (next_n is even here) + MAX_NUM_W_IN_REG = 40 if next_n >= 4 else 52 + NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads, self.weight_cache_limit) + if warpgroup_idx == 0: + # Math WG 0: process group 0 + tAcc_ref = tCtAcc_base_0[(None, None, None, 0)][((None, None), 0, 0)] + tAcc_ref_epi = cute.flat_divide(tAcc_ref, epi_sub_mn) + tiled_copy_ref = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_ref_epi[(None, None, 0, 0)] + ) + tTR_cC = tiled_copy_ref.get_slice(local_tidx).partition_D(cC) + m_coord = tTR_cC[0][0] + tTR_rAcc = cute.make_fragment_like(tTR_cC, self.acc_dtype) + tTR_rAcc_next = cute.make_fragment_like(tTR_cC, self.acc_dtype) + w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype) + q_stage_local = cutlass.Int32(0) + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + # Q pipeline consumer: wait for Q+Weights SMEM + if q_idx != q_idx_old: + if q_idx_old < batch_size: + q_pipeline.consumer_release(q_cons_state) + q_cons_state.advance() + q_pipeline.consumer_wait(q_cons_state) + q_stage_local = q_cons_state.index + # Preload first NUM_W_IN_REG weights per slot + for t_i in cutlass.range_constexpr(next_n): + for w_j in cutlass.range_constexpr(NUM_W_IN_REG): + w_cache[t_i * NUM_W_IN_REG + w_j] = sW[ + (t_i * num_heads + w_j, q_stage_local) + ] + + # Process KV block for group 0 (kv_idx + 0) + kv_pos = kv_idx * block_kv_val + m_coord + kv_pipeline_0.consumer_wait(kv_cons_state_math_0) + scale_val = sScales_0[(m_coord, kv_cons_state_math_0.index)] + for hh in cutlass.range_constexpr(self.n_halves): + umma_pipeline_0.consumer_wait(umma_cons_state_0) + if cutlass.const_expr(hh == self.n_halves - 1): + # The MMAs of both halves are complete: the KV stage is free. + kv_pipeline_0.consumer_release(kv_cons_state_math_0) + kv_cons_state_math_0.advance() + tAcc_c = tCtAcc_base_0[(None, None, None, umma_cons_state_0.index)][ + ((None, None), 0, 0) + ] + tAcc_c_epi = cute.flat_divide(tAcc_c, epi_sub_mn) + tc = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_c_epi[(None, None, 0, 0)] + ) + tTR = tc.get_slice(local_tidx).partition_S(tAcc_c_epi) + cute.copy(tc, tTR[(None, None, None, 0, 0)], tTR_rAcc) + cute.arch.fence_view_async_tmem_load() + for j in cutlass.range_constexpr(rows_per_half): + t = hh * rows_per_half + j + cur = tTR_rAcc if j % 2 == 0 else tTR_rAcc_next + if cutlass.const_expr(j + 1 < rows_per_half): + # Next row's TMEM load into the other buffer + cute.copy( + tc, + tTR[(None, None, None, 0, j + 1)], + tTR_rAcc_next if j % 2 == 0 else tTR_rAcc, + ) + if cutlass.const_expr(j == rows_per_half - 1): + # Every TMEM load of this half has completed: release its stage. + umma_pipeline_0.consumer_release(umma_cons_state_0) + result_t = self._row_sum( + cur.load(), t, w_cache, sW, q_stage_local, NUM_W_IN_REG + ) + out_row = q_idx * next_n + t + mLogits[(out_row, kv_pos)] = self.output_dtype(result_t * scale_val) + if cutlass.const_expr(j + 1 < rows_per_half): + cute.arch.fence_view_async_tmem_load() + umma_cons_state_0.advance() + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = ( + mContextLens[next_q_idx] + block_kv_val - 1 + ) // block_kv_val + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + # Release last Q stage (WG 0) + if q_idx < batch_size: + q_pipeline.consumer_release(q_cons_state) + q_cons_state.advance() + else: + # Math WG 1: process group 1 + tAcc_ref = tCtAcc_base_1[(None, None, None, 0)][((None, None), 0, 0)] + tAcc_ref_epi = cute.flat_divide(tAcc_ref, epi_sub_mn) + tiled_copy_ref = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_ref_epi[(None, None, 0, 0)] + ) + tTR_cC = tiled_copy_ref.get_slice(local_tidx).partition_D(cC) + m_coord = tTR_cC[0][0] + tTR_rAcc = cute.make_fragment_like(tTR_cC, self.acc_dtype) + tTR_rAcc_next = cute.make_fragment_like(tTR_cC, self.acc_dtype) + w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype) + q_stage_local = cutlass.Int32(0) + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + # Q pipeline consumer: wait for Q+Weights SMEM + if q_idx != q_idx_old: + if q_idx_old < batch_size: + q_pipeline.consumer_release(q_cons_state) + q_cons_state.advance() + q_pipeline.consumer_wait(q_cons_state) + q_stage_local = q_cons_state.index + # Preload first NUM_W_IN_REG weights per slot + for t_i in cutlass.range_constexpr(next_n): + for w_j in cutlass.range_constexpr(NUM_W_IN_REG): + w_cache[t_i * NUM_W_IN_REG + w_j] = sW[ + (t_i * num_heads + w_j, q_stage_local) + ] + + # Process KV block for group 1 (kv_idx + 1) + kv_pos = (kv_idx + 1) * block_kv_val + m_coord + kv_pipeline_1.consumer_wait(kv_cons_state_math_1) + scale_val = sScales_1[(m_coord, kv_cons_state_math_1.index)] + for hh in cutlass.range_constexpr(self.n_halves): + umma_pipeline_1.consumer_wait(umma_cons_state_1) + if cutlass.const_expr(hh == self.n_halves - 1): + # The MMAs of both halves are complete: the KV stage is free. + kv_pipeline_1.consumer_release(kv_cons_state_math_1) + kv_cons_state_math_1.advance() + tAcc_c = tCtAcc_base_1[(None, None, None, umma_cons_state_1.index)][ + ((None, None), 0, 0) + ] + tAcc_c_epi = cute.flat_divide(tAcc_c, epi_sub_mn) + tc = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_c_epi[(None, None, 0, 0)] + ) + tTR = tc.get_slice(local_tidx).partition_S(tAcc_c_epi) + cute.copy(tc, tTR[(None, None, None, 0, 0)], tTR_rAcc) + cute.arch.fence_view_async_tmem_load() + for j in cutlass.range_constexpr(rows_per_half): + t = hh * rows_per_half + j + cur = tTR_rAcc if j % 2 == 0 else tTR_rAcc_next + if cutlass.const_expr(j + 1 < rows_per_half): + # Next row's TMEM load into the other buffer + cute.copy( + tc, + tTR[(None, None, None, 0, j + 1)], + tTR_rAcc_next if j % 2 == 0 else tTR_rAcc, + ) + if cutlass.const_expr(j == rows_per_half - 1): + # Every TMEM load of this half has completed: release its stage. + umma_pipeline_1.consumer_release(umma_cons_state_1) + result_t = self._row_sum( + cur.load(), t, w_cache, sW, q_stage_local, NUM_W_IN_REG + ) + out_row = q_idx * next_n + t + mLogits[(out_row, kv_pos)] = self.output_dtype(result_t * scale_val) + if cutlass.const_expr(j + 1 < rows_per_half): + cute.arch.fence_view_async_tmem_load() + umma_cons_state_1.advance() + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = ( + mContextLens[next_q_idx] + block_kv_val - 1 + ) // block_kv_val + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + # Release last Q stage (WG 1) + if q_idx < batch_size: + q_pipeline.consumer_release(q_cons_state) + q_cons_state.advance() + else: + if warpgroup_idx == 0: + # Math WG 0: process group 0 + # Reference setup (stage 0) for m_coord + # flat_divide by sub-tile to get sub-tile partitions + tAcc_0_ref = tCtAcc_base_0[(None, None, None, 0)][((None, None), 0, 0)] + tAcc_0_ref_epi = cute.flat_divide(tAcc_0_ref, epi_sub_mn) + tiled_copy_ref_0 = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_0_ref_epi[(None, None, 0, 0)] + ) + thr_copy_ref_0 = tiled_copy_ref_0.get_slice(local_tidx) + tTR_cC = thr_copy_ref_0.partition_D(cC) + m_coord = tTR_cC[0][0] + + tTR_rAcc = cute.make_fragment_like(tTR_cC, self.acc_dtype) + + # Weight register cache: only first NUM_W_IN_REG + # per next_n slot (like DeepGEMM min(52, kNumHeads)). + # Remaining weights read from SMEM in epilogue. + # FP16 weights use half the regs, so we can fit + # all heads for next_n <= 3. + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + MAX_NUM_W_IN_REG = 64 if next_n <= 3 else 48 + else: + MAX_NUM_W_IN_REG = 64 if next_n == 1 else 40 if next_n >= 4 else 52 + NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads, self.weight_cache_limit) + w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype) + q_stage_local = cutlass.Int32(0) + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + # Q pipeline consumer: wait for Q+Weights SMEM + if q_idx != q_idx_old: + if q_idx_old < batch_size: + q_pipeline.consumer_release(q_cons_state) + q_cons_state.advance() + q_pipeline.consumer_wait(q_cons_state) + q_stage_local = q_cons_state.index + # Preload first NUM_W_IN_REG weights per slot + for t_i in cutlass.range_constexpr(next_n): + for w_j in cutlass.range_constexpr(NUM_W_IN_REG): + w_cache[t_i * NUM_W_IN_REG + w_j] = sW[ + (t_i * num_heads + w_j, q_stage_local) + ] + + # Process KV block for group 0 (kv_idx + 0) + # Unconditional Math (like DeepGEMM): OOB results + # written to aligned padding region in logits buffer. + kv_pos = kv_idx * block_kv_val + m_coord + + if cutlass.const_expr(self.remove_kv_wait_in_epilogue): + # Skip KV wait, rely on UMMA barrier's + # transitive visibility. + umma_pipeline_0.consumer_wait(umma_cons_state_0) + else: + # Default: wait KV first to overlap lds with + # UMMA computation. + kv_pipeline_0.consumer_wait(kv_cons_state_math_0) + sc_stage_0 = kv_cons_state_math_0.index + scale_val = sScales_0[(m_coord, sc_stage_0)] + umma_pipeline_0.consumer_wait(umma_cons_state_0) + + # --- TMEM sub-tile setup --- + # flat_divide accumulator by sub-tile shape; + # partition once, then loop over sub-tiles. + tCtAcc_c0 = tCtAcc_base_0[(None, None, None, umma_cons_state_0.index)] + tAcc_c0 = tCtAcc_c0[((None, None), 0, 0)] + tAcc_c0_epi = cute.flat_divide(tAcc_c0, epi_sub_mn) + tc_0 = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_c0_epi[(None, None, 0, 0)] + ) + tr_0 = tc_0.get_slice(local_tidx) + tTR_0 = tr_0.partition_S(tAcc_c0_epi) + + # --- First sub-tile LDTM + KV release --- + if cutlass.const_expr(self.early_tmem_copy): + # Issue first sub-tile LDTM early + cute.copy(tc_0, tTR_0[(None, None, None, 0, 0)], tTR_rAcc) + # Scale LDS + KV release fill latency + if cutlass.const_expr(self.remove_kv_wait_in_epilogue): + sc_stage_0 = kv_cons_state_math_0.index + scale_val = sScales_0[(m_coord, sc_stage_0)] + kv_pipeline_0.consumer_release(kv_cons_state_math_0) + kv_cons_state_math_0.advance() + cute.arch.fence_view_async_tmem_load() + else: + # Default: scale LDS + KV release first + if cutlass.const_expr(self.remove_kv_wait_in_epilogue): + sc_stage_0 = kv_cons_state_math_0.index + scale_val = sScales_0[(m_coord, sc_stage_0)] + kv_pipeline_0.consumer_release(kv_cons_state_math_0) + kv_cons_state_math_0.advance() + cute.copy(tc_0, tTR_0[(None, None, None, 0, 0)], tTR_rAcc) + cute.arch.fence_view_async_tmem_load() + + # --- Sub-tile compute loop --- + # Each sub-tile: LDTM.xN → fence → load → + # ReLU+FMA. Breaks FMA chain (16→4 per chunk) + # and interleaves LDTM with FP32 compute to + # reduce ShadowPipeThrottle. + # Sub-tiles are within each t-slot (num_heads + # // num_epi_subtiles wide). flat_divide yields + # next_n * num_epi_subtiles sub-tiles total; + # global index = t * num_epi_subtiles + i. + subtile_n = num_heads // num_epi_subtiles + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + packed_zero = pack_f16x2(Float16(0.0), Float16(0.0)) + for t in cutlass.range_constexpr(next_n): + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + ps0 = packed_zero + ps1 = packed_zero + else: + s0x = cutlass.Float32(0.0) + s0y = cutlass.Float32(0.0) + s1x = cutlass.Float32(0.0) + s1y = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(num_epi_subtiles): + # LDTM for sub-tiles 1..N-1 + # (sub-tile 0 handled above) + if t > 0 or i > 0: + cute.copy( + tc_0, + tTR_0[(None, None, None, 0, t * num_epi_subtiles + i)], + tTR_rAcc, + ) + cute.arch.fence_view_async_tmem_load() + # Release UMMA after last LDTM+fence + if t == next_n - 1 and i == num_epi_subtiles - 1: + umma_pipeline_0.consumer_release(umma_cons_state_0) + umma_cons_state_0.advance() + acc_vec = tTR_rAcc.load() + # Reg-path: weights from registers + reg_h_end = min(subtile_n, max(0, NUM_W_IN_REG - i * subtile_n)) + for h in cutlass.range_constexpr(0, reg_h_end, 4): + n0 = h + h_g = i * subtile_n + h + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + pa01 = pack_f16x2( + Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1]) + ) + pa23 = pack_f16x2( + Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3]) + ) + pa01 = max_f16x2(pa01, packed_zero) + pa23 = max_f16x2(pa23, packed_zero) + r0 = t * NUM_W_IN_REG + h_g + pw01 = pack_f16x2(w_cache[r0], w_cache[r0 + 1]) + pw23 = pack_f16x2(w_cache[r0 + 2], w_cache[r0 + 3]) + ps0 = fma_f16x2(pa01, pw01, ps0) + ps1 = fma_f16x2(pa23, pw23, ps1) + else: + a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0)) + a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0)) + a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0)) + a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0)) + r0 = t * NUM_W_IN_REG + h_g + w0 = w_cache[r0] + w1 = w_cache[r0 + 1] + w2 = w_cache[r0 + 2] + w3 = w_cache[r0 + 3] + s0x, s0y = cute.arch.fma_packed_f32x2( + (a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN + ) + s1x, s1y = cute.arch.fma_packed_f32x2( + (a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN + ) + # SMEM-path: weights from shared mem + smem_h_start = max(0, NUM_W_IN_REG - i * subtile_n) + for h in cutlass.range_constexpr(smem_h_start, subtile_n, 4): + n0 = h + h_g = i * subtile_n + h + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + pa01 = pack_f16x2( + Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1]) + ) + pa23 = pack_f16x2( + Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3]) + ) + pa01 = max_f16x2(pa01, packed_zero) + pa23 = max_f16x2(pa23, packed_zero) + pw01 = pack_f16x2( + sW[(t * num_heads + h_g, q_stage_local)], + sW[(t * num_heads + h_g + 1, q_stage_local)], + ) + pw23 = pack_f16x2( + sW[(t * num_heads + h_g + 2, q_stage_local)], + sW[(t * num_heads + h_g + 3, q_stage_local)], + ) + ps0 = fma_f16x2(pa01, pw01, ps0) + ps1 = fma_f16x2(pa23, pw23, ps1) + else: + a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0)) + a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0)) + a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0)) + a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0)) + w0 = sW[(t * num_heads + h_g, q_stage_local)] + w1 = sW[(t * num_heads + h_g + 1, q_stage_local)] + w2 = sW[(t * num_heads + h_g + 2, q_stage_local)] + w3 = sW[(t * num_heads + h_g + 3, q_stage_local)] + s0x, s0y = cute.arch.fma_packed_f32x2( + (a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN + ) + s1x, s1y = cute.arch.fma_packed_f32x2( + (a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN + ) + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + ps_sum = add_f16x2(ps0, ps1) + sum_lo, sum_hi = unpack_f16x2(ps_sum) + result_t = sum_lo + sum_hi + else: + result_t = s0x + s0y + s1x + s1y + out_row = q_idx * next_n + t + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + mLogits[(out_row, kv_pos)] = self.output_dtype( + result_t * Float16(scale_val) + ) + else: + mLogits[(out_row, kv_pos)] = self.output_dtype(result_t * scale_val) + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = ( + mContextLens[next_q_idx] + block_kv_val - 1 + ) // block_kv_val + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + # Release last Q stage (WG 0) + if q_idx < batch_size: + q_pipeline.consumer_release(q_cons_state) + q_cons_state.advance() + + else: + # Math WG 1: process group 1 + tAcc_1_ref = tCtAcc_base_1[(None, None, None, 0)][((None, None), 0, 0)] + tAcc_1_ref_epi = cute.flat_divide(tAcc_1_ref, epi_sub_mn) + tiled_copy_ref_1 = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_1_ref_epi[(None, None, 0, 0)] + ) + thr_copy_ref_1 = tiled_copy_ref_1.get_slice(local_tidx) + tTR_cC = thr_copy_ref_1.partition_D(cC) + m_coord = tTR_cC[0][0] + + tTR_rAcc = cute.make_fragment_like(tTR_cC, self.acc_dtype) + + # Weight register cache: only first NUM_W_IN_REG + # per next_n slot (like DeepGEMM min(52, kNumHeads)). + # FP16 weights use half the regs, so we can fit + # all heads for next_n <= 3. + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + MAX_NUM_W_IN_REG = 64 if next_n <= 3 else 48 + else: + MAX_NUM_W_IN_REG = 64 if next_n == 1 else 40 if next_n >= 4 else 52 + NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads, self.weight_cache_limit) + w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype) + q_stage_local = cutlass.Int32(0) + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + # Q pipeline consumer: wait for Q+Weights SMEM + if q_idx != q_idx_old: + if q_idx_old < batch_size: + q_pipeline.consumer_release(q_cons_state) + q_cons_state.advance() + q_pipeline.consumer_wait(q_cons_state) + q_stage_local = q_cons_state.index + # Preload first NUM_W_IN_REG weights per slot + for t_i in cutlass.range_constexpr(next_n): + for w_j in cutlass.range_constexpr(NUM_W_IN_REG): + w_cache[t_i * NUM_W_IN_REG + w_j] = sW[ + (t_i * num_heads + w_j, q_stage_local) + ] + + # Process KV block for group 1 (kv_idx + 1) + # Unconditional Math (like DeepGEMM) + kv_idx_1 = kv_idx + 1 + + kv_pos = kv_idx_1 * block_kv_val + m_coord + + if cutlass.const_expr(self.remove_kv_wait_in_epilogue): + umma_pipeline_1.consumer_wait(umma_cons_state_1) + else: + kv_pipeline_1.consumer_wait(kv_cons_state_math_1) + sc_stage_1 = kv_cons_state_math_1.index + scale_val = sScales_1[(m_coord, sc_stage_1)] + umma_pipeline_1.consumer_wait(umma_cons_state_1) + + # --- TMEM sub-tile setup (WG1) --- + tCtAcc_c1 = tCtAcc_base_1[(None, None, None, umma_cons_state_1.index)] + tAcc_c1 = tCtAcc_c1[((None, None), 0, 0)] + tAcc_c1_epi = cute.flat_divide(tAcc_c1, epi_sub_mn) + tc_1 = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_c1_epi[(None, None, 0, 0)] + ) + tr_1 = tc_1.get_slice(local_tidx) + tTR_1 = tr_1.partition_S(tAcc_c1_epi) + + # --- First sub-tile LDTM + KV release (WG1) --- + if cutlass.const_expr(self.early_tmem_copy): + cute.copy(tc_1, tTR_1[(None, None, None, 0, 0)], tTR_rAcc) + if cutlass.const_expr(self.remove_kv_wait_in_epilogue): + sc_stage_1 = kv_cons_state_math_1.index + scale_val = sScales_1[(m_coord, sc_stage_1)] + kv_pipeline_1.consumer_release(kv_cons_state_math_1) + kv_cons_state_math_1.advance() + cute.arch.fence_view_async_tmem_load() + else: + if cutlass.const_expr(self.remove_kv_wait_in_epilogue): + sc_stage_1 = kv_cons_state_math_1.index + scale_val = sScales_1[(m_coord, sc_stage_1)] + kv_pipeline_1.consumer_release(kv_cons_state_math_1) + kv_cons_state_math_1.advance() + cute.copy(tc_1, tTR_1[(None, None, None, 0, 0)], tTR_rAcc) + cute.arch.fence_view_async_tmem_load() + + # --- Sub-tile compute loop (WG1) --- + subtile_n = num_heads // num_epi_subtiles + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + packed_zero = pack_f16x2(Float16(0.0), Float16(0.0)) + for t in cutlass.range_constexpr(next_n): + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + ps0 = packed_zero + ps1 = packed_zero + else: + s0x = cutlass.Float32(0.0) + s0y = cutlass.Float32(0.0) + s1x = cutlass.Float32(0.0) + s1y = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(num_epi_subtiles): + if t > 0 or i > 0: + cute.copy( + tc_1, + tTR_1[(None, None, None, 0, t * num_epi_subtiles + i)], + tTR_rAcc, + ) + cute.arch.fence_view_async_tmem_load() + if t == next_n - 1 and i == num_epi_subtiles - 1: + umma_pipeline_1.consumer_release(umma_cons_state_1) + umma_cons_state_1.advance() + acc_vec = tTR_rAcc.load() + # Reg-path + reg_h_end = min(subtile_n, max(0, NUM_W_IN_REG - i * subtile_n)) + for h in cutlass.range_constexpr(0, reg_h_end, 4): + n0 = h + h_g = i * subtile_n + h + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + pa01 = pack_f16x2( + Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1]) + ) + pa23 = pack_f16x2( + Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3]) + ) + pa01 = max_f16x2(pa01, packed_zero) + pa23 = max_f16x2(pa23, packed_zero) + r0 = t * NUM_W_IN_REG + h_g + pw01 = pack_f16x2(w_cache[r0], w_cache[r0 + 1]) + pw23 = pack_f16x2(w_cache[r0 + 2], w_cache[r0 + 3]) + ps0 = fma_f16x2(pa01, pw01, ps0) + ps1 = fma_f16x2(pa23, pw23, ps1) + else: + a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0)) + a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0)) + a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0)) + a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0)) + r0 = t * NUM_W_IN_REG + h_g + w0 = w_cache[r0] + w1 = w_cache[r0 + 1] + w2 = w_cache[r0 + 2] + w3 = w_cache[r0 + 3] + s0x, s0y = cute.arch.fma_packed_f32x2( + (a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN + ) + s1x, s1y = cute.arch.fma_packed_f32x2( + (a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN + ) + # SMEM-path + smem_h_start = max(0, NUM_W_IN_REG - i * subtile_n) + for h in cutlass.range_constexpr(smem_h_start, subtile_n, 4): + n0 = h + h_g = i * subtile_n + h + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + pa01 = pack_f16x2( + Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1]) + ) + pa23 = pack_f16x2( + Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3]) + ) + pa01 = max_f16x2(pa01, packed_zero) + pa23 = max_f16x2(pa23, packed_zero) + pw01 = pack_f16x2( + sW[(t * num_heads + h_g, q_stage_local)], + sW[(t * num_heads + h_g + 1, q_stage_local)], + ) + pw23 = pack_f16x2( + sW[(t * num_heads + h_g + 2, q_stage_local)], + sW[(t * num_heads + h_g + 3, q_stage_local)], + ) + ps0 = fma_f16x2(pa01, pw01, ps0) + ps1 = fma_f16x2(pa23, pw23, ps1) + else: + a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0)) + a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0)) + a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0)) + a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0)) + w0 = sW[(t * num_heads + h_g, q_stage_local)] + w1 = sW[(t * num_heads + h_g + 1, q_stage_local)] + w2 = sW[(t * num_heads + h_g + 2, q_stage_local)] + w3 = sW[(t * num_heads + h_g + 3, q_stage_local)] + s0x, s0y = cute.arch.fma_packed_f32x2( + (a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN + ) + s1x, s1y = cute.arch.fma_packed_f32x2( + (a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN + ) + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + ps_sum = add_f16x2(ps0, ps1) + sum_lo, sum_hi = unpack_f16x2(ps_sum) + result_t = sum_lo + sum_hi + else: + result_t = s0x + s0y + s1x + s1y + out_row = q_idx * next_n + t + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + mLogits[(out_row, kv_pos)] = self.output_dtype( + result_t * Float16(scale_val) + ) + else: + mLogits[(out_row, kv_pos)] = self.output_dtype(result_t * scale_val) + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = ( + mContextLens[next_q_idx] + block_kv_val - 1 + ) // block_kv_val + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + # Release last Q stage (WG 1) + if q_idx < batch_size: + q_pipeline.consumer_release(q_cons_state) + q_cons_state.advance() + + # TMEM dealloc: math warps are allocator + last consumer + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + else: + cute.arch.warpgroup_reg_dealloc(24) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_hist_topk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_hist_topk.py new file mode 100644 index 000000000000..a9973f360be2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_hist_topk.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shard geometry and exchange workspace of the key-sharded decode indexer top-k. + +The decode indexer is replicated over the tensor-parallel group: without sharding every rank +computes the same FP8 MQA logits over the whole context and the same top-2048. Instead rank ``r`` +scores only its contiguous shard of the request's indexer keys (``glm_dsa_hist_shard_meta`` builds +the shard's context length and block table), and ``glm_dsa_hist_topk`` finds every row's exact +top-2048 over all shards by exchanging key histograms through a private multicast workspace; every +rank ends with the same rows. See ``cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.h``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Dict, Optional + +import torch + +from tensorrt_llm.mapping import Mapping + +# Constants of glmDsaHistTopK.h. +MAX_RANKS = 16 +SHARD_ALIGN = 128 +MAX_SHARD = 98304 +MAX_ROWS = 8 +TOPK = 2048 +WORKSPACE_WORDS = 3 * (3 * 8 * 16 * 1088 + 8 * 2048 + 8 * 4096 * 2) +STATE_WORDS = 4 + 8 * 4 + + +def shard_size(kv_len: int, group_size: int) -> int: + """Keys per rank of a request of ``kv_len`` indexer keys: ceil(kv_len / group_size), rounded + up to SHARD_ALIGN. Rank ``r`` holds the keys [r * S, r * S + n_r), n_r = clamp(kv_len - r * S, + 0, S).""" + per_rank = -(-max(kv_len, 0) // group_size) + return -(-per_rank // SHARD_ALIGN) * SHARD_ALIGN + + +def shard_capacity(max_seq_len: int, group_size: int) -> int: + """Largest shard of a request of at most ``max_seq_len`` indexer keys.""" + return shard_size(max_seq_len, group_size) + + +@dataclass +class HistTopkWorkspace: + """The exchange workspace and device state of ``glm_dsa_hist_topk`` for one TP group. + + ``comm_buffer`` is this rank's unicast view of ``WORKSPACE_WORDS`` int32 words of a + McastGPUBuffer, all -1 at creation; ``state`` is ``STATE_WORDS`` int32 words, zero at creation. + Both are private to the op and live as long as the engine: captured CUDA graphs launch on + them. ``handle`` keeps the multicast allocation alive. + """ + + group_size: int + handle: Any + comm_buffer: torch.Tensor + state: torch.Tensor + comm: Any + + +def hist_topk_supported(mapping: Mapping) -> bool: + """Whether the key-sharded top-k can serve this mapping: a pure tensor-parallel group of 2 to + 16 ranks without attention DP or context parallelism, on an MNNVL fabric.""" + from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce + + return ( + 2 <= mapping.tp_size <= MAX_RANKS + and mapping.world_size == mapping.tp_size + and not mapping.has_cp() + and not mapping.enable_attention_dp + and MNNVLAllReduce.is_mnnvl(mapping, torch.bfloat16) + ) + + +def create_hist_topk_workspace(mapping: Mapping, comm: Optional[Any] = None) -> HistTopkWorkspace: + """Allocate the workspace on every rank of ``mapping``'s TP group (a collective call). + + ``comm`` is the TP-group communicator the multicast handles are exchanged over: the MNNVL + all-reduce's (an MPI split, or the TP ProcessGroup without MPI) when None. The outcome is + converged over the group, so every rank raises or none does. + """ + from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, + _make_mnnvl_mcast_buffer, + _mnnvl_device_index, + _mnnvl_workspace_all_succeeded, + ) + + if comm is None: + comm = _get_mnnvl_workspace_comm(mapping) + error: Optional[Exception] = None + workspace = None + try: + handle = _make_mnnvl_mcast_buffer( + comm, WORKSPACE_WORDS * 4, mapping, use_fabric_handle=mapping.is_multi_node() + ) + comm_buffer = handle.get_uc_buffer(mapping.tp_rank, (WORKSPACE_WORDS,), torch.int32, 0) + with torch.inference_mode(): + comm_buffer.fill_(-1) + state = torch.zeros( + STATE_WORDS, + dtype=torch.int32, + device=torch.device("cuda", _mnnvl_device_index(mapping)), + ) + torch.cuda.synchronize() + workspace = HistTopkWorkspace(mapping.tp_size, handle, comm_buffer, state, comm) + except Exception as e: + # Any failure is converged over the group and re-raised on every rank below: a rank that + # raised alone would leave its peers waiting in the next exchange. + error = e + if not _mnnvl_workspace_all_succeeded(comm, error is None): + raise RuntimeError( + "GLM DSA histogram top-k workspace construction failed on at least one rank" + ) from error + return workspace + + +_WORKSPACES: Dict[Mapping, HistTopkWorkspace] = {} + + +def get_hist_topk_workspace(mapping: Mapping) -> HistTopkWorkspace: + """The workspace of ``mapping``'s TP group, created on the first call (collective: every rank + of the group makes its first call in the same place, outside CUDA graph capture).""" + workspace = _WORKSPACES.get(mapping) + if workspace is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("the GLM DSA histogram top-k workspace is created outside capture") + workspace = _WORKSPACES[mapping] = create_hist_topk_workspace(mapping) + return workspace diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_indexer_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_indexer_prologue.py new file mode 100644 index 000000000000..5ecdcf2bc7f0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_indexer_prologue.py @@ -0,0 +1,470 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused DSA indexer prologue: one Triton launch for the indexer K and Q of a decode step. + +``_dsa_indexer_prologue_fused_kernel`` computes in one launch what the indexer's per-kernel chain +computes between the indexer projections and the paged MQA logits: + + 1. ``k_norm`` LayerNorm over the indexer head dim; + 2. RoPE on the Q and K rope slices (the flashinfer ``apply_rope_with_cos_sin_cache_inplace`` + expression); + 3. ``[pe | nope]`` FP8 E4M3 quantization of Q and K with one scale per 128 values + (``fused_cat_fp8``'s rule); + 4. ``weights * q_scale * softmax_scale * n_heads**-0.5``; + 5. the scatter of each K row (128 FP8 bytes and a 4-byte scale) into the indexer K cache + (``indexer_k_cache_scatter_op``'s layout and slot rule). + +Each step has the numerics of the chain's kernel: fp32 math on bf16-rounded intermediates (the +LayerNorm output and the RoPE output are rounded to bf16 where the chain materialises bf16 +tensors), ``fused_cat_fp8``'s scale (UE8M0 power-of-two round-up or plain amax / 448), its +``rcp.approx.ftz`` reciprocal and RN-satfinite E4M3 conversion, and the scatter kernel's flat-index +to strided-offset unravel. + +``launch_indexer_prologue`` runs the whole prologue. ``launch_indexer_k_tail`` runs the same kernel +with the Q part compiled out (``WRITE_Q=False``), one program per token: the K part of a step whose +indexer wq_b epilogue computes the Q part. + +With ``use_pdl`` the kernel is launched with programmatic dependent launch and waits on +``griddepcontrol.wait`` before it reads its inputs: + +- ``launch_dependents``: each program triggers the dependent launch right after its wait, so the + next kernel on the stream can start while this one runs. Only for a dependent that reads this + kernel's outputs (Q, the weights, the new K-cache rows) after its own wait and writes nothing + before it: the early-KV decode MQA-logits kernel. +- ``prewait``: each program loads the position, the cos/sin row, the K LayerNorm weight and bias + and the slot mappings before its wait, and Q, K and the weights right after it. Only when the + positions and slot mappings were written before the kernels this one depends on started (the + step's metadata preparation). +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _rcp_approx_ftz(x): + # fused_cat_fp8 uses the MUFU reciprocal (rcp.approx.ftz.f32), not IEEE 1/x. + return tl.inline_asm_elementwise( + "rcp.approx.ftz.f32 $0, $1;", + "=r,r", + [x], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _fp8_row_scale(amax, UE8M0: tl.constexpr): + """fused_cat_fp8 per-row scale: amax clamped at 1e-12, / 448, optional UE8M0 round-up.""" + amax = tl.maximum(amax, 1.0e-12) + ratio = amax * 0.002232142857142857 # 1.0f / 448.0f + if UE8M0: + bits = ratio.to(tl.uint32, bitcast=True) + mant = bits & 0x007FFFFF + exp_bits = bits & 0x7F800000 + exp_bits = tl.where(mant != 0, exp_bits + 0x00800000, exp_bits) + scale = exp_bits.to(tl.float32, bitcast=True) + else: + scale = ratio + return scale + + +@triton.jit +def _rope_terms(d, ROPE_DIM: tl.constexpr, INTERLEAVE: tl.constexpr): + """Partner column, cos/sin column and sign for column ``d`` (flashinfer layout).""" + HALF: tl.constexpr = ROPE_DIM // 2 + is_rope = d < ROPE_DIM + if INTERLEAVE: + # GPT-J style: pairs (2i, 2i + 1) share cos/sin column i; even lanes subtract. + partner = tl.where(is_rope, d ^ 1, d) + cs_col = d // 2 + neg = (d % 2) == 0 + else: + # NeoX style: halves [0, HALF) and [HALF, ROPE_DIM); the first half subtracts. + partner = tl.where(d < HALF, d + HALF, tl.where(is_rope, d - HALF, d)) + cs_col = d % HALF + neg = d < HALF + cs_col = tl.where(is_rope, cs_col, 0) + return is_rope, partner, cs_col, neg + + +@triton.jit +def _apply_rope(x, xp, c, s, neg): + # flashinfer: vec[i] = vec[i] * cos[i] + (neg ? -perm[i] : perm[i]) * sin[i]; + # nvcc contracts this to fma(vec, cos, (+-perm) * sin). + t = tl.where(neg, -xp, xp) * s + return tl.fma(x, c, t) + + +@triton.jit +def _unravel_offset(flat, d1, d2, d3, s0, s1, s2, s3): + """indexerKCacheScatter's flatIndexToMemoryOffset over a [d0, d1, d2, d3] strided view.""" + i3 = flat % d3 + flat = flat // d3 + i2 = flat % d2 + flat = flat // d2 + i1 = flat % d1 + i0 = flat // d1 + return i0 * s0 + i1 * s1 + i2 * s2 + i3 * s3 + + +@triton.jit +def _dsa_indexer_prologue_fused_kernel( + # inputs + k_ptr, # [N, D] pre-norm indexer K (bf16 or fp32; rounded to bf16 first) + k_stride, + ln_w_ptr, # [D] LayerNorm weight + ln_b_ptr, # [D] LayerNorm bias + q_ptr, # [N, H, D] bf16 (wq_b output view, unit inner stride) + q_stride_tok, + q_stride_head, + w_ptr, # [N, H] weights_proj output (bf16 or fp32) + w_stride, + pos_ptr, # [N] positions + cos_sin_ptr, # [max_pos, ROPE_DIM] fp32: [cos(0..HALF) | sin(0..HALF)] + cos_sin_stride, + # outputs + q_fp8_ptr, # [N, H, D] fp8 e4m3 (contiguous) + q_scale_ptr, # [N, H] fp32 + k_fp8_ptr, # [N, D] fp8 e4m3 (contiguous) + k_scale_ptr, # [N] fp32 + w_out_ptr, # [N, H] fp32 + # indexer K cache (uint8 [d0, d1, d2, d3] strided view) + slot mappings + k_cache_ptr, + slot_fp8_ptr, + slot_scale_ptr, + n_write, + cache_d1, + cache_d2, + cache_d3, + cache_s0, + cache_s1, + cache_s2, + cache_s3, + # scalars + eps, + weight_scale, + H: tl.constexpr, + D: tl.constexpr, + ROPE_DIM: tl.constexpr, + HEADS_PER_PROG: tl.constexpr, + INTERLEAVE: tl.constexpr, + UE8M0: tl.constexpr, + WRITE_CACHE: tl.constexpr, + USE_PDL: tl.constexpr, + LAUNCH_DEPENDENTS: tl.constexpr, + PREWAIT: tl.constexpr, + WRITE_Q: tl.constexpr, +): + tok = tl.program_id(0) + hblk = tl.program_id(1) + HALF: tl.constexpr = ROPE_DIM // 2 + if PREWAIT: + # Positions, cos/sin, the K LayerNorm and the slot mappings: written before the + # kernels this one depends on run. + tok64 = tok.to(tl.int64) + d = tl.arange(0, D) + is_rope, partner, cs_col, neg = _rope_terms(d, ROPE_DIM, INTERLEAVE) + k_prog = hblk == 0 + pos = tl.load(pos_ptr + tok).to(tl.int64) + cs_row = cos_sin_ptr + pos * cos_sin_stride + cos_v = tl.load(cs_row + cs_col, mask=is_rope, other=1.0) + sin_v = tl.load(cs_row + HALF + cs_col, mask=is_rope, other=0.0) + lw = tl.load(ln_w_ptr + d, mask=k_prog, other=0.0).to(tl.float32) + lb = tl.load(ln_b_ptr + d, mask=k_prog, other=0.0).to(tl.float32) + lwp = tl.load(ln_w_ptr + partner, mask=k_prog, other=0.0).to(tl.float32) + lbp = tl.load(ln_b_ptr + partner, mask=k_prog, other=0.0).to(tl.float32) + if WRITE_CACHE: + slot = tl.load(slot_fp8_ptr + tok, mask=k_prog & (tok < n_write), other=-1) + slot_s = tl.load(slot_scale_ptr + tok, mask=k_prog & (tok < n_write), other=-1) + if USE_PDL: + tl.extra.cuda.gdc_wait() + if LAUNCH_DEPENDENTS: + tl.extra.cuda.gdc_launch_dependents() + + if not PREWAIT: + tok64 = tok.to(tl.int64) + d = tl.arange(0, D) + is_rope, partner, cs_col, neg = _rope_terms(d, ROPE_DIM, INTERLEAVE) + + pos = tl.load(pos_ptr + tok).to(tl.int64) + cs_row = cos_sin_ptr + pos * cos_sin_stride + cos_v = tl.load(cs_row + cs_col, mask=is_rope, other=1.0) + sin_v = tl.load(cs_row + HALF + cs_col, mask=is_rope, other=0.0) + + if WRITE_Q: + # ---------------- Q: HEADS_PER_PROG heads of this token ---------------- + h = hblk * HEADS_PER_PROG + tl.arange(0, HEADS_PER_PROG) + q_row = q_ptr + tok64 * q_stride_tok + h[:, None].to(tl.int64) * q_stride_head + qx = tl.load(q_row + d[None, :]).to(tl.float32) + qp = tl.load(q_row + partner[None, :]).to(tl.float32) + if PREWAIT: + # Every load from the kernels this one waits for, before the first store. + w = tl.load(w_ptr + tok64 * w_stride + h).to(tl.float32) + k_row = k_ptr + tok64 * k_stride + kx = tl.load(k_row + d, mask=k_prog, other=0.0).to(tl.bfloat16).to(tl.float32) + kxp = tl.load(k_row + partner, mask=k_prog, other=0.0).to(tl.bfloat16).to(tl.float32) + q_rot = _apply_rope(qx, qp, cos_v[None, :], sin_v[None, :], neg[None, :]) + qv = tl.where(is_rope[None, :], q_rot, qx).to(tl.bfloat16).to(tl.float32) + q_amax = tl.max(tl.abs(qv), axis=1) + q_scale = _fp8_row_scale(q_amax, UE8M0) + q_inv = _rcp_approx_ftz(q_scale) + q8 = (qv * q_inv[:, None]).to(tl.float8e4nv) + qo = (tok64 * H + h[:, None]) * D + d[None, :] + tl.store(q_fp8_ptr + qo, q8) + tl.store(q_scale_ptr + tok64 * H + h, q_scale) + + # weights * q_scale * (softmax_scale * n_heads^-0.5), same op order as _scale(). + if not PREWAIT: + w = tl.load(w_ptr + tok64 * w_stride + h).to(tl.float32) + tl.store(w_out_ptr + tok64 * H + h, (w * q_scale) * weight_scale) + elif PREWAIT: + # K only: the K row, right after the wait. + k_row = k_ptr + tok64 * k_stride + kx = tl.load(k_row + d, mask=k_prog, other=0.0).to(tl.bfloat16).to(tl.float32) + kxp = tl.load(k_row + partner, mask=k_prog, other=0.0).to(tl.bfloat16).to(tl.float32) + + # ---------------- K: once per token ---------------- + if hblk == 0: + if not PREWAIT: + k_row = k_ptr + tok64 * k_stride + kx = tl.load(k_row + d).to(tl.bfloat16).to(tl.float32) + kxp = tl.load(k_row + partner).to(tl.bfloat16).to(tl.float32) + mean = tl.sum(kx, axis=0) / D + kc = kx - mean + var = tl.sum(kc * kc, axis=0) / D + rstd = tl.math.rsqrt(var + eps) + if not PREWAIT: + lw = tl.load(ln_w_ptr + d).to(tl.float32) + lb = tl.load(ln_b_ptr + d).to(tl.float32) + lwp = tl.load(ln_w_ptr + partner).to(tl.float32) + lbp = tl.load(ln_b_ptr + partner).to(tl.float32) + kn = (kc * rstd * lw + lb).to(tl.bfloat16).to(tl.float32) + knp = ((kxp - mean) * rstd * lwp + lbp).to(tl.bfloat16).to(tl.float32) + k_rot = _apply_rope(kn, knp, cos_v, sin_v, neg) + kv = tl.where(is_rope, k_rot, kn).to(tl.bfloat16).to(tl.float32) + k_amax = tl.max(tl.abs(kv), axis=0) + k_scale = _fp8_row_scale(k_amax, UE8M0) + # inline asm needs a block operand: broadcast the row scale over D. + k_inv = _rcp_approx_ftz(k_scale + tl.zeros([D], dtype=tl.float32)) + k8 = (kv * k_inv).to(tl.float8e4nv) + tl.store(k_fp8_ptr + tok64 * D + d, k8) + tl.store(k_scale_ptr + tok64, k_scale) + + if WRITE_CACHE: + if tok < n_write: + if not PREWAIT: + slot = tl.load(slot_fp8_ptr + tok) + slot_s = tl.load(slot_scale_ptr + tok) + if (slot >= 0) & (slot_s >= 0): + # 4-byte chunks, each chunk base unravelled (as the scatter kernel). + chunk = slot + (d // 4) * 4 + dst = _unravel_offset( + chunk, cache_d1, cache_d2, cache_d3, cache_s0, cache_s1, cache_s2, cache_s3 + ) + (d % 4) + tl.store(k_cache_ptr + dst, k8.to(tl.uint8, bitcast=True)) + b = tl.arange(0, 4) + sbits = k_scale.to(tl.int32, bitcast=True) + sbytes = ((sbits >> (b * 8)) & 0xFF).to(tl.uint8) + dst_s = _unravel_offset( + slot_s, cache_d1, cache_d2, cache_d3, cache_s0, cache_s1, cache_s2, cache_s3 + ) + tl.store(k_cache_ptr + dst_s + b, sbytes) + + +def _heads_per_prog(n_heads: int) -> int: + for hp in (8, 4, 2, 1): + if n_heads % hp == 0: + return hp + return 1 + + +def _cache_args( + k_cache: torch.Tensor, + slot_mapping_fp8: torch.Tensor, + slot_mapping_scale: torch.Tensor, + num_cache_tokens: int, +) -> tuple: + """The K-cache operands: the uint8 [d0, d1, d2, d3] view, the slot mappings, the row count, + and the view's inner sizes and strides.""" + cache_u8 = k_cache.view(torch.uint8) + cd = cache_u8.shape + cs = cache_u8.stride() + return ( + cache_u8, + slot_mapping_fp8, + slot_mapping_scale, + int(num_cache_tokens), + cd[1], + cd[2], + cd[3], + cs[0], + cs[1], + cs[2], + cs[3], + ) + + +def launch_indexer_prologue( + indexer_k: torch.Tensor, + weights: torch.Tensor, + q: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + eps: float, + weight_scale: float, + rope_dim: int, + interleave: bool, + use_ue8m0: bool, + k_cache: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + num_cache_tokens: int, + use_pdl: bool, + launch_dependents: bool, + prewait: bool, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """The whole prologue in one launch. + + Args: + indexer_k: [N, D] pre-norm indexer K (bf16, or fp32 which is rounded to bf16), unit inner + stride. + weights: [N, H] weights_proj output (bf16 or fp32), unit inner stride. + q: [N, H, D] bf16 indexer Q (the wq_b output), unit inner stride. + positions: [N] (any int dtype) rope positions. + cos_sin: [max_pos, rope_dim] fp32 cos/sin cache, ``[cos | sin]`` per row. + k_cache / slot_mapping_*: when ``k_cache`` is given and ``num_cache_tokens`` > 0, the first + ``num_cache_tokens`` K rows are also written into the indexer K cache (the layout and + slot rule of ``indexer_k_cache_scatter_op``). + use_pdl / launch_dependents / prewait: see the module docstring. + + Returns (q_fp8 [N, H, D], k_fp8 [N, D], k_scale [N, 1], weights [N, H] fp32, + q_scale [N, H, 1]). + """ + n_tok, n_heads, head_dim = q.shape + dev = q.device + q_fp8 = torch.empty((n_tok, n_heads, head_dim), dtype=torch.float8_e4m3fn, device=dev) + q_scale = torch.empty((n_tok, n_heads, 1), dtype=torch.float32, device=dev) + k_fp8 = torch.empty((n_tok, head_dim), dtype=torch.float8_e4m3fn, device=dev) + k_scale = torch.empty((n_tok, 1), dtype=torch.float32, device=dev) + w_out = torch.empty((n_tok, n_heads), dtype=torch.float32, device=dev) + if n_tok == 0: + return q_fp8, k_fp8, k_scale, w_out, q_scale + + write_cache = k_cache is not None and num_cache_tokens > 0 + if write_cache: + cache_args = _cache_args(k_cache, slot_mapping_fp8, slot_mapping_scale, num_cache_tokens) + else: + cache_args = (k_fp8, k_scale, k_scale, 0, 1, 1, 1, 0, 0, 0, 0) + + hp = _heads_per_prog(n_heads) + _dsa_indexer_prologue_fused_kernel[(n_tok, n_heads // hp)]( + indexer_k, + indexer_k.stride(0), + ln_weight, + ln_bias, + q, + q.stride(0), + q.stride(1), + weights, + weights.stride(0), + positions, + cos_sin, + cos_sin.stride(0), + q_fp8, + q_scale, + k_fp8, + k_scale, + w_out, + *cache_args, + float(eps), + float(weight_scale), + H=n_heads, + D=head_dim, + ROPE_DIM=rope_dim, + HEADS_PER_PROG=hp, + INTERLEAVE=bool(interleave), + UE8M0=bool(use_ue8m0), + WRITE_CACHE=write_cache, + USE_PDL=use_pdl, + LAUNCH_DEPENDENTS=use_pdl and launch_dependents, + PREWAIT=use_pdl and prewait, + WRITE_Q=True, + num_warps=4, + launch_pdl=use_pdl, + ) + return q_fp8, k_fp8, k_scale, w_out, q_scale + + +def launch_indexer_k_tail( + indexer_k: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + eps: float, + rope_dim: int, + interleave: bool, + use_ue8m0: bool, + k_fp8: torch.Tensor, + k_scale: torch.Tensor, + k_cache: torch.Tensor, + slot_mapping_fp8: torch.Tensor, + slot_mapping_scale: torch.Tensor, + num_cache_tokens: int, + use_pdl: bool, + launch_dependents: bool, + prewait: bool, +) -> None: + """The K part of ``launch_indexer_prologue``, one program per token: the K LayerNorm, RoPE and + FP8 quantization of K into ``k_fp8`` [N, D] / ``k_scale`` [N, 1], and the first + ``num_cache_tokens`` rows into the indexer K cache. The kernel and its K arithmetic are those of + the whole prologue; only its Q part is compiled out (H and HEADS_PER_PROG select nothing there). + """ + n_tok = indexer_k.shape[0] + if n_tok == 0: + return + _dsa_indexer_prologue_fused_kernel[(n_tok, 1)]( + indexer_k, + indexer_k.stride(0), + ln_weight, + ln_bias, + # The Q operands are not read without WRITE_Q. + indexer_k, + 0, + 0, + indexer_k, + 0, + positions, + cos_sin, + cos_sin.stride(0), + k_fp8, + k_scale, + k_fp8, + k_scale, + k_scale, + *_cache_args(k_cache, slot_mapping_fp8, slot_mapping_scale, num_cache_tokens), + float(eps), + 1.0, + H=1, + D=indexer_k.shape[1], + ROPE_DIM=rope_dim, + HEADS_PER_PROG=1, + INTERLEAVE=bool(interleave), + UE8M0=bool(use_ue8m0), + WRITE_CACHE=True, + USE_PDL=use_pdl, + LAUNCH_DEPENDENTS=use_pdl and launch_dependents, + PREWAIT=use_pdl and prewait, + WRITE_Q=False, + num_warps=4, + launch_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_indexer_reference.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_indexer_reference.py new file mode 100644 index 000000000000..a35396b535d0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_indexer_reference.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Plain-torch references of the indexer prologue's arithmetic, shared by its two entries. + +fp32 math rounded where the kernel rounds: the LayerNorm output and the RoPE output to bf16, one +rounding for each fused multiply-add, ``fused_cat_fp8``'s per-row scale and E4M3 conversion, and +the indexer K-cache scatter's byte placement. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch + + +def fma(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> torch.Tensor: + """fp32 ``a * b + c`` with one rounding: the fp64 product of two fp32 values is exact.""" + return (a.double() * b.double() + c.double()).float() + + +def rope_rows( + x: torch.Tensor, positions: torch.Tensor, cos_sin: torch.Tensor, rope_dim: int, interleave: bool +) -> torch.Tensor: + """RoPE on the first ``rope_dim`` values of the fp32 rows ``x`` [..., D] at ``positions`` + (one per leading row), rounded to bf16 and returned as fp32. + + ``cos_sin`` rows are ``[cos | sin]`` of ``rope_dim / 2`` values each. Interleaved pairs + (2i, 2i + 1) share column i and the even value subtracts; otherwise the two halves pair up and + the first half subtracts. + """ + d = torch.arange(x.shape[-1], device=x.device) + half = rope_dim // 2 + is_rope = d < rope_dim + if interleave: + partner = torch.where(is_rope, d ^ 1, d) + cs_col = d // 2 + neg = d % 2 == 0 + else: + partner = torch.where(d < half, d + half, torch.where(is_rope, d - half, d)) + cs_col = d % half + neg = d < half + cs_col = torch.where(is_rope, cs_col, torch.zeros_like(cs_col)) + table = cos_sin[positions.long()] + while table.dim() < x.dim(): + table = table.unsqueeze(-2) + cos = torch.where(is_rope, table[..., cs_col], torch.ones_like(table[..., cs_col])) + sin = torch.where(is_rope, table[..., half + cs_col], torch.zeros_like(table[..., cs_col])) + xp = x[..., partner] + rotated = fma(x, cos, torch.where(neg, -xp, xp) * sin) + return torch.where(is_rope, rotated, x).to(torch.bfloat16).float() + + +def quantize_rows(v: torch.Tensor, use_ue8m0: bool) -> Tuple[torch.Tensor, torch.Tensor]: + """``fused_cat_fp8``'s rule over the last dim of fp32 ``v``: (E4M3 values, fp32 scales with a + trailing dim of 1). scale = max(amax, 1e-12) / 448, rounded up to a power of two with + ``use_ue8m0``; values = e4m3(v * (1 / scale)).""" + amax = v.abs().amax(dim=-1, keepdim=True).clamp_min(1.0e-12) + ratio = amax * 0.002232142857142857 + if use_ue8m0: + bits = ratio.view(torch.int32) + exp_bits = bits & 0x7F800000 + exp_bits = torch.where((bits & 0x007FFFFF) != 0, exp_bits + 0x00800000, exp_bits) + scale = exp_bits.view(torch.float32) + else: + scale = ratio + # The kernel multiplies by the MUFU reciprocal, exact for a power-of-two scale. + return (v * (1.0 / scale)).to(torch.float8_e4m3fn), scale + + +def reference_k_rows( + indexer_k: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + eps: float, + rope_dim: int, + interleave: bool, + use_ue8m0: bool, +) -> Tuple[torch.Tensor, torch.Tensor]: + """``(k_fp8 [N, D], k_scale [N, 1])``: LayerNorm of bf16(indexer_k) rounded to bf16, RoPE + rounded to bf16, then the E4M3 rule.""" + rows = indexer_k.shape[0] + kx = indexer_k.to(torch.bfloat16).float() + mean = kx.mean(dim=1, keepdim=True) + kc = kx - mean + rstd = torch.rsqrt((kc * kc).mean(dim=1, keepdim=True) + eps) + kn = fma(kc * rstd, ln_weight.float()[None, :], ln_bias.float()[None, :]).to(torch.bfloat16) + kv = rope_rows(kn.float(), positions[:rows], cos_sin, rope_dim, interleave) + return quantize_rows(kv, use_ue8m0) + + +def reference_q_rows( + q: torch.Tensor, + weights: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + weight_scale: float, + rope_dim: int, + interleave: bool, + use_ue8m0: bool, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``(q_fp8 [N, H, D], q_scale [N, H, 1], weights [N, H] fp32)``: RoPE of the bf16 Q rounded + to bf16, the E4M3 rule per head, and ``(weights * q_scale) * weight_scale`` in fp32.""" + rows = q.shape[0] + qv = rope_rows(q.float(), positions[:rows], cos_sin, rope_dim, interleave) + q8, q_scale = quantize_rows(qv, use_ue8m0) + w = (weights.float() * q_scale[..., 0]) * torch.tensor(weight_scale, dtype=torch.float32) + return q8, q_scale, w + + +def unravel(flat: torch.Tensor, shape: torch.Size) -> Tuple[torch.Tensor, ...]: + """The 4-D index of each logical flat index of a [d0, d1, d2, d3] view.""" + d1, d2, d3 = shape[1:] + i3 = flat % d3 + flat = flat // d3 + i2 = flat % d2 + flat = flat // d2 + i1 = flat % d1 + return flat // d1, i1, i2, i3 + + +def write_cache_rows( + cache: torch.Tensor, + k8: torch.Tensor, + k_scale: torch.Tensor, + slot_mapping_fp8: torch.Tensor, + slot_mapping_scale: torch.Tensor, + num_cache_tokens: int, +) -> None: + """The indexer K-cache scatter, in place: for each of the first ``num_cache_tokens`` rows with + both slots >= 0, its 128 FP8 bytes at the logical flat byte index ``slot_mapping_fp8[row]`` of + the 4-D byte view ``cache`` and its fp32 scale's 4 little-endian bytes at + ``slot_mapping_scale[row]``.""" + view = cache.view(torch.uint8) + fp8_slots = slot_mapping_fp8[:num_cache_tokens].long() + scale_slots = slot_mapping_scale[:num_cache_tokens].long() + keep = (fp8_slots >= 0) & (scale_slots >= 0) + rows = torch.nonzero(keep).flatten() + if rows.numel() == 0: + return + data = k8.view(torch.uint8)[rows] + scale_bytes = k_scale.reshape(-1, 1).contiguous().view(torch.uint8)[rows] + width = data.shape[1] + fp8_index = fp8_slots[rows, None] + torch.arange(width, device=view.device) + scale_index = scale_slots[rows, None] + torch.arange(4, device=view.device) + view[unravel(fp8_index, view.shape)] = data + view[unravel(scale_index, view.shape)] = scale_bytes diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_mtp_topk_stash.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_mtp_topk_stash.py new file mode 100644 index 000000000000..4e41c60fedaa --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_mtp_topk_stash.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MTP index share: each request's last accepted TopK row of draft step 0, in one Triton launch. + +With MTP index share, draft step 0's indexer computes the TopK of the verify step's next_n rows, +and the draft steps after it reuse, per request, the row of its last accepted token. The kernel +copies those rows, the rows ``torch.index_select(gen_topk, 0, index)`` would copy for +``index[r] = r * next_n + clamp(num_accepted[r] - 1, 0, next_n - 1)``. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _stash_last_accepted_rows_kernel( + src_ptr, # [num_requests * next_n, n_cols] int32 TopK rows + src_stride, + num_accepted_ptr, # [num_requests] accepted tokens per request + dst_ptr, # [num_requests, n_cols] int32 + dst_stride, + next_n, + n_cols, + BLOCK: tl.constexpr, +): + req = tl.program_id(0) + chunk = tl.program_id(1) + # The request's last accepted row of its next_n: num_accepted - 1, clamped to [0, next_n - 1]. + accepted = tl.load(num_accepted_ptr + req) + offset = tl.minimum(tl.maximum(accepted - 1, 0), next_n - 1) + row = req.to(tl.int64) * next_n + offset.to(tl.int64) + cols = chunk * BLOCK + tl.arange(0, BLOCK) + mask = cols < n_cols + values = tl.load(src_ptr + row * src_stride + cols, mask=mask) + tl.store(dst_ptr + req.to(tl.int64) * dst_stride + cols, values, mask=mask) + + +def launch_stash_last_accepted_rows( + gen_topk: torch.Tensor, num_accepted: torch.Tensor, dst: torch.Tensor, next_n: int +) -> None: + """dst[r] = gen_topk[r * next_n + clamp(num_accepted[r] - 1, 0, next_n - 1)] for each request. + + Args: + gen_topk: [num_requests * next_n, n_cols] int32 TopK rows, unit column stride. + num_accepted: [num_requests] accepted tokens per request (contiguous). + dst: [num_requests, n_cols] int32, unit column stride. + """ + num_requests, n_cols = dst.shape + if num_requests == 0: + return + block = min(triton.next_power_of_2(n_cols), 2048) + _stash_last_accepted_rows_kernel[(num_requests, triton.cdiv(n_cols, block))]( + gen_topk, + gen_topk.stride(0), + num_accepted, + dst, + dst.stride(0), + next_n, + n_cols, + BLOCK=block, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/convert_req_index_to_global.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/convert_req_index_to_global.py new file mode 100644 index 000000000000..eabfcfedf014 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/convert_req_index_to_global.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Request-local TopK token indices to rows of the MLA KV pool, via the trtllm convert op.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +def _spec(rows, **overrides) -> dict: + spec = dict( + rows=rows, + topk=2048, + tokens_per_block=64, + kv_len=530_000, + table_width=8_288, + pool_pages=9_000, + page_index_scale=79, + layer=78, + invalid=0, + table_layout="rows", + ) + spec.update(overrides) + return spec + + +class GlmConvertReqIndexToGlobal(OpWrapper): + """``out[t, c]`` = the KV-pool row of request-local token ``tok = token_indices[t, c]`` of + request ``req_id[t]``: ``block_table[req, tok // block_size] * stride_factor + layer_id * + block_size + tok % block_size``; -1 where ``tok < 0``, ``tok // block_size`` is past the table + row, or the table entry is negative. Columns ``[0, num_topk_tokens)``; returns a fresh tensor + shaped like ``token_indices``. + + For the MLA pool of layer ``l`` the target passes ``stride_factor = page_index_scale * + tokens_per_block`` and ``layer_id = layer_offset`` from the KV cache manager's + ``get_primary_pool_page_index_params(l)``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="c1 draft step 1: the reused TopK row of one request, MTP layer 78", + spec=_spec(1), + ), + Cell( + why="-1 TopK slots of a short request stay -1", + spec=_spec(1, kv_len=1_500, invalid=548), + ), + Cell( + why="four requests, table rows of the generation requests only", + spec=_spec(4, kv_len=70_000, table_width=1_200, table_layout="gen_rows"), + ), + ) + + note = """ + Metadata consumed: none. At c1 draft step 1 the target passes the request map of the draft + step (``req_idx_per_token`` of the generation tokens minus the context count), the generation + requests' rows of the MLA block table, the stashed TopK rows ([num_requests, 2048] int32), + the page size, 2048, and the MTP layer's stride factor and layer offset. A plain launch. + + The op reads ``token_indices`` and the table through ``.contiguous()`` copies when they are + strided. + """ + + def __call__( + self, + req_id: torch.Tensor, + block_table: torch.Tensor, + token_indices: torch.Tensor, + block_size: int, + num_topk_tokens: int, + stride_factor: int, + layer_id: int, + ) -> torch.Tensor: + return torch.ops.trtllm.convert_req_index_to_global( + req_id, block_table, token_indices, block_size, num_topk_tokens, stride_factor, layer_id + ) + + def reference( + self, + req_id: torch.Tensor, + block_table: torch.Tensor, + token_indices: torch.Tensor, + block_size: int, + num_topk_tokens: int, + stride_factor: int, + layer_id: int, + ) -> torch.Tensor: + """The pool rows in int64 arithmetic, cast back to int32.""" + tok = token_indices[:, :num_topk_tokens].long() + page = tok.clamp_min(0) // block_size + valid_page = page < block_table.shape[1] + req = req_id.long()[:, None].expand_as(tok) + base = block_table[req, page.clamp(max=block_table.shape[1] - 1)].long() + rows = base * stride_factor + tok.clamp_min(0) % block_size + layer_id * block_size + valid = (tok >= 0) & valid_page & (base >= 0) + out = token_indices.clone() + out[:, :num_topk_tokens] = torch.where(valid, rows, torch.full_like(rows, -1)).int() + return out + + def is_valid( + self, + req_id: torch.Tensor, + block_table: torch.Tensor, + token_indices: torch.Tensor, + block_size: int, + num_topk_tokens: int, + stride_factor: int, + layer_id: int, + ) -> None: + # Columns past num_topk_tokens are left unwritten in the fresh output. + assert num_topk_tokens == token_indices.shape[1], ( + f"num_topk_tokens={num_topk_tokens} leaves columns of the " + f"{token_indices.shape[1]}-wide output unwritten" + ) + assert req_id.shape[0] == token_indices.shape[0], "one request index per TopK row" + # A request index past the table reads another allocation's memory. + assert int(req_id.max()) < block_table.shape[0] and int(req_id.min()) >= 0, ( + "a request index past the block table" + ) + # The pool row is computed in int32. + largest = int(block_table.max()) * stride_factor + (layer_id + 1) * block_size + assert largest < 2**31, "pool rows overflow int32" + + +convert_req_index_to_global = GlmConvertReqIndexToGlobal() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/fp8_paged_mqa_logits.py new file mode 100644 index 000000000000..6ef531814cdb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/fp8_paged_mqa_logits.py @@ -0,0 +1,250 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Indexer FP8 MQA logits over paged keys (CuTe DSL), the whole verify window as one query atom.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._fp8_paged_mqa_logits import SPLIT_KV, num_sms, run_fp8_paged_mqa_logits + +#: The indexer K cache's page: 64 tokens of 128 FP8 bytes, then their 64 fp32 scales. +_PAGE = 64 +#: Columns of the c1 shard logits: the largest shard of the c1 gen server's indexer capacity +#: (1,048,596 keys) on 16 ranks. +_SHARD_CAPACITY = 65_664 + + +def _spec(context_lens, max_context_len=_SHARD_CAPACITY, **overrides) -> dict: + spec = dict( + next_n=8, + num_heads=32, + head_dim=128, + tokens_per_block=_PAGE, + context_lens=tuple(context_lens), + max_context_len=max_context_len, + early_kv=True, + data="int", + pool_layout="contiguous", + ) + spec.update(overrides) + return spec + + +class GlmFp8PagedMqaLogits(OpWrapper): + """``logits[b * next_n + t, j] = k_scale[j] * sum_h weights[b * next_n + t, h] * + relu(q[b, t, h] . k[j])`` for the keys ``j < context_lens[b]`` of request ``b``. + + The keys live in the indexer K cache: key ``j`` of request ``b`` is token ``j % page`` of page + ``block_table[b, j // page]``, whose bytes hold the page's FP8 rows and then their fp32 scales. + ``weights`` already carry the Q scales and the softmax scale (the indexer prologue's or the + wq_b epilogue's output). ``schedule_meta`` partitions the requests' 256-key splits over the + device's SMs (``paged_mqa_logits_schedule``); the kernel runs one persistent CTA per SM. + + A request's ``next_n`` query rows are one MMA atom of ``next_n * num_heads`` columns. With + ``early_kv`` (an even ``next_n`` whose half-atom is a whole number of 16-column UMMA atoms) the + N-split variant runs, launched with PDL: it loads the KV blocks that end before each request's + newest ``next_n`` keys before ``griddepcontrol.wait`` and everything else after it. Both + variants compute every element in the same order. + + ``logits`` is the caller's [B * next_n, >= round_up(max_context_len, 256)] fp32 buffer: the + kernel writes whole 256-key splits, so columns from ``context_lens[b]`` to the split's end hold + unspecified values. Returns the view ``logits[:, :max_context_len]``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="c1 key shard, 8 rows x 32 heads, early KV: a full 33,152-key shard", + spec=_spec([33_152]), + ), + Cell( + why="the request's last shard: 32,720 keys, a partial final page", + spec=_spec([32_720]), + ), + Cell( + why="an empty shard scores one key", + spec=_spec([1]), + ), + Cell( + why="the default kernel (early KV off) on the c1 shard", + spec=_spec([33_152], early_kv=False), + ), + Cell( + why="UE8M0-scaled keys and real-valued weights (accumulation order visible)", + spec=_spec([20_000], data="randn"), + ), + Cell( + why="the whole context on every rank (no key sharding): 200,000 keys of 1M", + spec=_spec([200_000], max_context_len=1_048_576), + ), + Cell( + why="cache pages interleaved with other layers' pages (strided page dim)", + spec=_spec([5_000], pool_layout="interleaved"), + ), + Cell( + why="four requests of 8 rows each (a larger batch's verify step)", + spec=_spec([3_000, 70_000, 129, 40_000], max_context_len=131_072), + ), + ) + + note = """ + Metadata consumed: none; every input is an argument. At c1 the target passes the step's Q + ([1, 8, 32, 128] E4M3) and scaled weights ([8, 32] fp32), the layer's indexer K-cache view + ([pages, 64, 1, 132] uint8), and, with the key-sharded top-k, the shard's context length, + block table and schedule (``hist_shard_meta``, ``paged_mqa_logits_schedule``) with + ``max_context_len`` = the shard capacity; without it, the request's indexer KV length, its + INDEX_KEY block-table row and the full schedule. + + The compiled kernel is cached per (page size, heads, head dim, next_n, SMs, weight-register + cache, early KV); the first call of a shape compiles (seconds), so it must run before CUDA + graph capture. + + PDL: only the early-KV variant is launched with the PDL attribute. It requires that the + kernels it waits for write no key older than each request's newest ``next_n`` and no block + table, context length or schedule entry. + + Not certified: FP16 epilogues, next_n other than 8, head counts other than 32, pages other + than 64 tokens. + """ + + def __call__( + self, + q: torch.Tensor, + kv_cache: torch.Tensor, + weights: torch.Tensor, + context_lens: torch.Tensor, + block_table: torch.Tensor, + schedule_meta: torch.Tensor, + logits: torch.Tensor, + max_context_len: int, + early_kv: bool, + ) -> torch.Tensor: + return run_fp8_paged_mqa_logits( + q, + kv_cache, + weights, + context_lens, + block_table, + schedule_meta, + logits, + max_context_len, + early_kv, + ) + + def reference( + self, + q: torch.Tensor, + kv_cache: torch.Tensor, + weights: torch.Tensor, + context_lens: torch.Tensor, + block_table: torch.Tensor, + schedule_meta: torch.Tensor, + logits: torch.Tensor, + max_context_len: int, + early_kv: bool, + ) -> torch.Tensor: + """[B * next_n, max_context_len] fp32, fp64-accumulated and rounded once; columns past a + request's context are 0. The schedule, the output buffer and the variant change where and + when the values are computed, not what they are.""" + batch, next_n, num_heads, head_dim = q.shape + page = kv_cache.shape[1] + out = torch.zeros(batch * next_n, max_context_len, dtype=torch.float32, device=q.device) + page_bytes = kv_cache.reshape(kv_cache.shape[0], -1) + for b, ctx in enumerate(context_lens.tolist()): + num_pages = -(-ctx // page) + pages = page_bytes[block_table[b, :num_pages].long()] + keys = pages[:, : page * head_dim].contiguous().view(torch.float8_e4m3fn) + keys = keys.reshape(num_pages * page, head_dim)[:ctx].float().double() + scales = pages[:, page * head_dim :].contiguous().view(torch.float32) + scales = scales.reshape(num_pages * page)[:ctx].double() + qb = q[b].reshape(next_n * num_heads, head_dim).float().double() + wb = weights[b * next_n : (b + 1) * next_n].reshape(next_n * num_heads, 1).double() + for start in range(0, ctx, 16_384): + stop = min(start + 16_384, ctx) + scores = torch.relu(qb @ keys[start:stop].t()) * wb + summed = scores.reshape(next_n, num_heads, stop - start).sum(dim=1) + out[b * next_n : (b + 1) * next_n, start:stop] = ( + summed * scales[start:stop] + ).float() + return out + + def is_valid( + self, + q: torch.Tensor, + kv_cache: torch.Tensor, + weights: torch.Tensor, + context_lens: torch.Tensor, + block_table: torch.Tensor, + schedule_meta: torch.Tensor, + logits: torch.Tensor, + max_context_len: int, + early_kv: bool, + ) -> None: + assert q.dim() == 4 and q.dtype == torch.float8_e4m3fn, ( + f"q must be [B, next_n, H, D] E4M3, got {tuple(q.shape)} {q.dtype}" + ) + batch, next_n, num_heads, head_dim = q.shape + # The kernel takes Q and the weights through compact layouts; a strided view would be read + # as if it were packed. + assert q.is_contiguous(), "q must be contiguous" + assert head_dim == 128, f"head_dim={head_dim}" + assert weights.shape == (batch * next_n, num_heads) and weights.is_contiguous(), ( + f"weights must be a contiguous [B * next_n, H] = [{batch * next_n}, {num_heads}] " + f"tensor, got {tuple(weights.shape)}" + ) + assert weights.dtype == torch.float32, weights.dtype + # A page is read as [page * D FP8 bytes][page * 4 scale bytes]: any other inner layout + # reads keys and scales from the wrong bytes. + assert kv_cache.dim() == 4 and kv_cache.dtype == torch.uint8, ( + "kv_cache must be the [pages, page, 1, D + 4] uint8 view of the indexer K cache" + ) + page = kv_cache.shape[1] + assert page in (32, 64, 128) and kv_cache.shape[2:] == (1, head_dim + 4), ( + f"kv_cache pages must be [{page}, 1, {head_dim + 4}] with page in (32, 64, 128)" + ) + assert kv_cache[0].is_contiguous(), "the bytes of one page must be contiguous" + assert context_lens.shape == (batch,) and context_lens.dtype == torch.int32, ( + "context_lens must be [B] int32" + ) + assert block_table.dim() == 2 and block_table.shape[0] == batch, ( + "block_table must hold one row per request" + ) + assert block_table.dtype == torch.int32 and block_table.is_contiguous(), ( + "block_table must be contiguous int32 rows" + ) + lens = context_lens.tolist() + longest = max(lens) if lens else 0 + # The kernel reads whole 128-key compute blocks of block-table entries. + needed = -(-longest // 128) * (128 // page) + assert block_table.shape[1] >= needed, ( + f"block_table rows of {block_table.shape[1]} pages cannot cover {longest} keys " + f"({needed} page entries are read)" + ) + if needed: + assert int(block_table[:, :needed].max()) < kv_cache.shape[0], ( + "a block-table entry points past the cache" + ) + assert longest <= max_context_len, ( + f"context length {longest} exceeds max_context_len={max_context_len}" + ) + sms = num_sms(q.device.index if q.device.index is not None else 0) + assert schedule_meta.shape == (sms + 1, 2) and schedule_meta.dtype == torch.int32, ( + f"schedule_meta must be the [{sms + 1}, 2] int32 schedule of this device's {sms} SMs" + ) + assert logits.dim() == 2 and logits.dtype == torch.float32 and logits.stride(1) == 1, ( + "logits must be an fp32 matrix with a unit column stride" + ) + assert logits.shape[0] == batch * next_n, ( + f"logits must hold B * next_n = {batch * next_n} rows, got {logits.shape[0]}" + ) + assert logits.shape[1] >= -(-max_context_len // SPLIT_KV) * SPLIT_KV, ( + f"logits rows of {logits.shape[1]} columns cannot take the whole {SPLIT_KV}-key " + f"splits of max_context_len={max_context_len}" + ) + + +fp8_paged_mqa_logits = GlmFp8PagedMqaLogits() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/hist_shard_meta.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/hist_shard_meta.py new file mode 100644 index 000000000000..f497903c8adf --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/hist_shard_meta.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""This rank's key shard of the generation request: its context length and block table.""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._hist_topk import shard_capacity, shard_size + +#: c1: 16 ranks, 64-token pages, and the shard table of the gen server's indexer capacity. +_C1 = dict(group_size=16, tokens_per_block=64, capacity=shard_capacity(1_048_596, 16)) + + +def _spec(kv_len, rank, **overrides) -> dict: + spec = dict(_C1, kv_len=kv_len, rank=rank) + spec.update(overrides) + return spec + + +class GlmHistShardMeta(OpWrapper): + """``shard_ctx[0]`` and ``shard_block_table[0, :]`` of rank ``rank`` for one request. + + With ``S = shard_size(kv_lens[0], group_size)`` (ceil(kv_len / group_size) rounded up to 128), + the rank holds the keys [rank * S, rank * S + n), ``n = clamp(kv_len - rank * S, 0, S)``: + + - ``shard_ctx[0] = max(n, 1)`` (an empty shard scores one key no row reaches); + - ``shard_block_table[0, j] = block_table[rank * S / tokens_per_block + j]`` for the shard's + ``ceil(n / tokens_per_block)`` pages, 0 for the rest of the row. + + ``block_table`` is the request's INDEX_KEY block-table row (one-dimensional). One CTA. Returns + ``(shard_ctx, shard_block_table)``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="c1: a 530,000-key request on rank 0", spec=_spec(530_000, 0)), + Cell(why="c1: the same request's last rank, a partial shard", spec=_spec(530_000, 15)), + Cell(why="the full 1,048,596-key capacity on rank 15", spec=_spec(1_048_596, 15)), + Cell(why="a short request: rank 13 holds no key, one unread key", spec=_spec(3_000, 13)), + Cell( + why="a shard ending inside a page: 1,000,003 keys on rank 15", + spec=_spec(1_000_003, 15), + ), + ) + + note = """ + Metadata consumed: none. The target passes the generation request's indexer KV length (int32, + on the device), its row of the INDEX_KEY block offsets, persistent ``shard_ctx`` [1] and + ``shard_block_table`` [1, ceil(capacity / 64)] int32 buffers (captured CUDA graphs read + them; capacity = ``_hist_topk.shard_capacity(indexer max_seq_len, 16)``), the group size + (16), this rank's TP rank and the page size (64). + + The kernel refuses a group above 16 ranks and a page size that does not divide 128. + """ + + def __call__( + self, + kv_lens: torch.Tensor, + block_table: torch.Tensor, + shard_ctx: torch.Tensor, + shard_block_table: torch.Tensor, + group_size: int, + rank: int, + tokens_per_block: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + torch.ops.trtllm.glm_dsa_hist_shard_meta( + kv_lens, block_table, shard_ctx, shard_block_table, group_size, rank, tokens_per_block + ) + return shard_ctx, shard_block_table + + def reference( + self, + kv_lens: torch.Tensor, + block_table: torch.Tensor, + shard_ctx: torch.Tensor, + shard_block_table: torch.Tensor, + group_size: int, + rank: int, + tokens_per_block: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Fresh ``(shard_ctx, shard_block_table)`` computed on the host.""" + kv_len = int(kv_lens[0]) + size = shard_size(kv_len, group_size) + base = rank * size + n = min(max(kv_len - base, 0), size) + first_page = base // tokens_per_block + pages = -(-n // tokens_per_block) + row = block_table.tolist() + width = shard_block_table.numel() + table = [ + row[first_page + j] if j < pages and first_page + j < len(row) else 0 + for j in range(width) + ] + ctx = torch.tensor([max(n, 1)], dtype=torch.int32, device=shard_ctx.device) + out = torch.tensor(table, dtype=torch.int32, device=shard_block_table.device) + return ctx, out.view(shard_block_table.shape) + + def is_valid( + self, + kv_lens: torch.Tensor, + block_table: torch.Tensor, + shard_ctx: torch.Tensor, + shard_block_table: torch.Tensor, + group_size: int, + rank: int, + tokens_per_block: int, + ) -> None: + # The kernel reads the first row of a two-dimensional table: another request's pages. + assert block_table.dim() == 1, "block_table must be the request's one-dimensional row" + # A table narrower than the shard's pages drops pages the logits then read as page 0. + kv_len = int(kv_lens[0]) + size = shard_size(kv_len, group_size) + n = min(max(kv_len - rank * size, 0), size) + assert shard_block_table.numel() * tokens_per_block >= n, ( + f"shard_block_table holds {shard_block_table.numel()} pages, the shard {n} keys" + ) + # Pages past the request's row are read as 0 instead of failing. + if n > 0: + first_page = rank * size // tokens_per_block + assert first_page + -(-n // tokens_per_block) <= block_table.numel(), ( + "the shard's pages run past the request's block-table row" + ) + + +hist_shard_meta = GlmHistShardMeta() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/hist_topk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/hist_topk.py new file mode 100644 index 000000000000..3051f1d2f81f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/hist_topk.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Exact top-2048 of the key-sharded decode indexer logits, by an MNNVL histogram exchange.""" + +from __future__ import annotations + +from typing import Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._hist_topk import MAX_SHARD, TOPK, shard_size + +_TOKEN_BITS = 21 + + +def _spec(group_size, kv_len, pattern="randn", rows=8, next_n=8) -> dict: + return dict(group_size=group_size, kv_len=kv_len, pattern=pattern, rows=rows, next_n=next_n) + + +def canonical_rows(rows: torch.Tensor) -> torch.Tensor: + """Each row's tokens in ascending order, the -1 slots last: the set a row holds, as a tensor.""" + largest = torch.iinfo(rows.dtype).max + keyed = torch.where(rows < 0, torch.full_like(rows, largest), rows) + ordered = keyed.sort(dim=1).values + return torch.where(ordered == largest, torch.full_like(ordered, -1), ordered) + + +def ordered_keys(logits: torch.Tensor) -> torch.Tensor: + """The kernel's order-preserving 32-bit key of each FP32 logit, as int64: the bits with the + sign bit set when it is clear and all bits inverted when it is set, so -0 sorts below +0, + positive NaNs above +inf and negative NaNs below -inf; the one key equal to the empty word + (NaN 0x7FFFFFFF) is clamped one below it.""" + bits = logits.contiguous().view(torch.int32).long() & 0xFFFFFFFF + keys = torch.where(bits >= 0x80000000, (~bits) & 0xFFFFFFFF, bits | 0x80000000) + return keys.clamp(max=0xFFFFFFFE) + + +class GlmHistTopk(OpWrapper): + """Each verify row's exact top-2048 over the TP group's key shards; every rank gets the same + rows. + + The request of ``kv_lens[0]`` indexer keys is split into contiguous shards of ``S = + shard_size(kv_len, group_size)`` keys, rank ``r`` holding [r * S, r * S + n_r); ``logits`` + [num_rows, >= S] fp32 holds this rank's shard (its column ``c`` is key r * S + c). Verify row + ``i`` of ``next_n`` sees the keys [0, kv_len - next_n + i + 1). + + ``out[i]`` [2048] int32 gets the row's top-2048 keys by (logit, key), the larger key first among + equal logits, in an unspecified order that is the same on every rank; a row of at most 2048 + keys gets [0, n) and then -1. FP32 values are ordered by their bits: -0 below +0, NaNs above + +inf (positive) or below -inf (negative). + + The exchange runs through ``comm_buffer`` (this rank's unicast view of the group's multicast + workspace) and ``state`` (``_hist_topk.HistTopkWorkspace``), both private to the op: every rank + of the group makes the same sequence of calls. Returns ``out``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="16 ranks (c1): a 530,000-key request, 8 verify rows", spec=_spec(16, 530_000)), + Cell( + why="16 ranks: the c1 gen server's 1,048,596-key capacity", + spec=_spec(16, 1_048_596), + ), + Cell( + why="16 ranks: the largest shard, 98,304 keys of a 1,572,864-key request", + spec=_spec(16, 16 * MAX_SHARD), + ), + Cell(why="4 ranks: a 300,000-key request, random logits", spec=_spec(4, 300_000)), + Cell(why="4 ranks: ties at the cut, the larger keys taken", spec=_spec(4, 200_000, "ties")), + Cell(why="4 ranks: every logit equal", spec=_spec(4, 50_000, "equal")), + Cell(why="4 ranks: +-inf, +-0 and NaN among the logits", spec=_spec(4, 120_000, "special")), + Cell( + why="4 ranks: one histogram bin holds more than 65,535 keys of a shard", + spec=_spec(4, 4 * MAX_SHARD, "saturated"), + ), + Cell( + why="4 ranks: 2,051 keys, rows 0-4 without an exchange, rows 5-7 with", + spec=_spec(4, 2_051), + ), + Cell(why="4 ranks: no row exceeds 2,048 keys", spec=_spec(4, 1_000)), + Cell(why="4 ranks: one row of a next_n = 1 call", spec=_spec(4, 90_000, rows=1, next_n=1)), + ) + + note = """ + Metadata consumed: none. At c1 the target passes the shard logits of + ``fp8_paged_mqa_logits`` (the view ``logits[:, :shard capacity]``), the generation request's + indexer KV length ([1] int32 on the device), its [8, 2048] TopK buffer, the workspace's + ``comm_buffer`` and ``state``, next_n = 8 and the TP size (16). The workspace is created once + per process (``_hist_topk.get_hist_topk_workspace``) by every rank of the group, outside CUDA + graph capture. + + Order and ties: the set is exact; the order within a row after a cut on the top 11 key bits + follows shared-memory atomics and can change between calls with the same logits. + + The launch is PDL-enabled (``TRTLLM_ENABLE_PDL``): before its wait it initializes its cluster + barriers and reads ``kv_lens`` (written by the step's metadata preparation, before the logits). + A wait for a peer longer than 120 s records the error in ``state`` and aborts the kernel. + + The kernel refuses (loudly) a shard above 98,304 keys, more than 8 rows or 16 ranks, logits rows + not 16-byte aligned, and a workspace of another group. + """ + + def __call__( + self, + logits: torch.Tensor, + kv_lens: torch.Tensor, + out: torch.Tensor, + comm_buffer: torch.Tensor, + state: torch.Tensor, + next_n: int, + group_size: int, + ) -> torch.Tensor: + torch.ops.trtllm.glm_dsa_hist_topk( + logits, kv_lens, out, comm_buffer, state, next_n, group_size + ) + return out + + def reference( + self, + logits: Sequence[torch.Tensor], + kv_lens: torch.Tensor, + out: torch.Tensor, + comm_buffer: torch.Tensor, + state: torch.Tensor, + next_n: int, + group_size: int, + ) -> torch.Tensor: + """``canonical_rows`` of the expected ``out``, from the group's shard logits in rank order. + The workspace only carries the exchange.""" + assert len(logits) == group_size, "the reference takes every rank's shard logits" + kv_len = int(kv_lens[0]) + size = shard_size(kv_len, group_size) + rows = out.shape[0] + expected = torch.full((rows, TOPK), -1, dtype=torch.int32, device=out.device) + whole = [] + for rank, shard in enumerate(logits): + n = min(max(kv_len - rank * size, 0), size) + whole.append(shard[:rows, :n].to(out.device)) + full = torch.cat(whole, dim=1) + for row in range(rows): + window = kv_len - next_n + row + 1 + if window <= TOPK: + count = max(window, 0) + expected[row, :count] = torch.arange(count, dtype=torch.int32, device=out.device) + continue + tokens = torch.arange(window, dtype=torch.int64, device=out.device) + composite = (ordered_keys(full[row, :window]) << _TOKEN_BITS) | tokens + top = composite.topk(TOPK).values & ((1 << _TOKEN_BITS) - 1) + expected[row] = top.sort().values.int() + return expected + + def is_valid( + self, + logits: torch.Tensor, + kv_lens: torch.Tensor, + out: torch.Tensor, + comm_buffer: torch.Tensor, + state: torch.Tensor, + next_n: int, + group_size: int, + ) -> None: + # The kernel reads kv_lens[0] only: a batch's lengths would select request 0's silently. + assert kv_lens.numel() == 1, ( + "kv_lens must hold the one generation request's indexer KV length" + ) + + +hist_topk = GlmHistTopk() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/indexer_k_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/indexer_k_tail.py new file mode 100644 index 000000000000..0bd627f19c12 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/indexer_k_tail.py @@ -0,0 +1,277 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Indexer K tail: k_norm, RoPE, E4M3 K and the indexer K-cache write in one Triton launch.""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._indexer_prologue import launch_indexer_k_tail +from ._indexer_reference import reference_k_rows, write_cache_rows + +#: GLM-5.2's indexer: head dim, RoPE slice (first 64 values, interleaved pairs), k_norm eps, the +#: UE8M0 scale format, and the indexer K cache's page (64 tokens of 128 FP8 bytes, then their 64 +#: fp32 scales). The pre-norm K is a column view of the fused-A output, whose rows are the +#: kv_a_proj_with_mqa (2624), indexer wk (128) and weights_proj (32) rows. +_GLM = dict( + head_dim=128, + rope_dim=64, + interleave=True, + use_ue8m0=True, + eps=1e-6, + tokens_per_block=64, + fused_a_width=2784, + k_col=2624, +) + + +def _spec(rows, first_position, **overrides) -> dict: + """``rows`` K rows of one request at positions first_position, first_position + 1, ...""" + spec = dict( + _GLM, + rows=rows, + cache_rows=rows, + first_position=first_position, + k_layout="fused_a", + pool_layout="contiguous", + pages=16, + negative_slot=False, + use_pdl=True, + launch_dependents=True, + prewait=True, + ) + spec.update(overrides) + return spec + + +class GlmIndexerKTail(OpWrapper): + """The K part of the fused indexer prologue for rows whose Q part another kernel computes. + + For every row ``i`` of ``indexer_k`` [N, 128] (the pre-norm indexer K): + + 1. ``kn = LayerNorm(bf16(indexer_k[i]); ln_weight, ln_bias, eps)``, fp32 math, rounded to bf16; + 2. RoPE on ``kn[:rope_dim]`` at ``positions[i]`` with the ``[cos | sin]`` row of ``cos_sin`` + (interleaved pairs when ``interleave``), ``fma(x, cos, (+-partner) * sin)``, rounded to bf16; + 3. ``k_scale[i] = max(amax, 1e-12) / 448``, rounded up to a power of two with ``use_ue8m0``; + ``k_fp8[i] = e4m3(k * rcp(k_scale[i]))``; + 4. for ``i < num_cache_tokens`` with both slots >= 0, the 128 FP8 bytes into + ``k_cache`` at the flat byte offset ``slot_mapping_fp8[i]`` and the 4 scale bytes at + ``slot_mapping_scale[i]``, each 4-byte chunk unravelled over ``k_cache``'s strided view + (``indexer_k_cache_scatter_op``'s rule). + + Returns ``(k_fp8, k_scale, k_cache)``, all written in place. + + With ``use_pdl`` the kernel waits on ``griddepcontrol.wait`` before reading its inputs; + ``launch_dependents`` lets the next kernel of the stream launch right after that wait (only for + the early-KV decode MQA logits, which read the new cache rows after their own wait); ``prewait`` + loads the positions, cos/sin rows, the LayerNorm parameters and the slot mappings before the + wait (only when the step's metadata preparation wrote them before the kernels this one waits + for). + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="verify / draft step 0: 8 rows from fused-A columns, PDL trigger + pre-wait", + spec=_spec(8, 529_981), + ), + Cell( + why="8 rows at the end of a 1M-token context (cos/sin rows near 2^20)", + spec=_spec(8, 1_048_560), + ), + Cell( + why="plain launch, PDL off", + spec=_spec(8, 4_093, use_pdl=False, launch_dependents=False, prewait=False), + ), + Cell( + why="PDL trigger without the pre-wait loads", + spec=_spec(8, 70_000, prewait=False), + ), + Cell( + why="padded graph rows: 8 K rows, the first 5 written to the cache", + spec=_spec(8, 130, cache_rows=5), + ), + Cell( + why="a slot of -1 leaves that row out of the cache", + spec=_spec(8, 640, negative_slot=True), + ), + Cell( + why="cache pages interleaved with other layers' pages (strided page dim)", + spec=_spec(8, 63, pool_layout="interleaved"), + ), + Cell( + why="contiguous pre-norm K (a separate wk|weights_proj GEMM)", + spec=_spec(8, 9_000, k_layout="contiguous"), + ), + ) + + note = """ + Metadata consumed: none; every input is an argument. The target passes the K columns + [2624, 2752) of the fused-A output, the step's ``position_ids`` (int32), its cos/sin cache + ``[max_pos, 64]`` fp32 (``[cos | sin]`` per row), the layer's ``k_norm`` weight and bias, the + layer's indexer K-cache view ``[pages, 64, 1, 132]`` uint8 and the step's + ``slot_mapping_fp8`` / ``slot_mapping_scale`` (int64 flat byte offsets), with + ``num_cache_tokens`` = the step's token count. + + ``k_fp8`` / ``k_scale`` are written in full (N rows) even where only the cache is read + afterwards. A slot < 0 skips the row's cache write, as the scatter op does. + + Not certified: fp32 pre-norm K, NeoX RoPE (``interleave=False``), the plain amax / 448 scale + (``use_ue8m0=False``), other head or RoPE dims, more than 8 rows. + """ + + def __call__( + self, + indexer_k: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + eps: float, + rope_dim: int, + interleave: bool, + use_ue8m0: bool, + k_fp8: torch.Tensor, + k_scale: torch.Tensor, + k_cache: torch.Tensor, + slot_mapping_fp8: torch.Tensor, + slot_mapping_scale: torch.Tensor, + num_cache_tokens: int, + use_pdl: bool, + launch_dependents: bool, + prewait: bool, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + launch_indexer_k_tail( + indexer_k, + positions, + cos_sin, + ln_weight, + ln_bias, + eps, + rope_dim, + interleave, + use_ue8m0, + k_fp8, + k_scale, + k_cache, + slot_mapping_fp8, + slot_mapping_scale, + num_cache_tokens, + use_pdl, + launch_dependents, + prewait, + ) + return k_fp8, k_scale, k_cache + + def reference( + self, + indexer_k: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + eps: float, + rope_dim: int, + interleave: bool, + use_ue8m0: bool, + k_fp8: torch.Tensor, + k_scale: torch.Tensor, + k_cache: torch.Tensor, + slot_mapping_fp8: torch.Tensor, + slot_mapping_scale: torch.Tensor, + num_cache_tokens: int, + use_pdl: bool, + launch_dependents: bool, + prewait: bool, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Fresh ``(k_fp8, k_scale, k_cache)`` from plain torch, fp32 math rounded where the kernel + rounds. PDL changes when the kernel runs, not what it computes, so the three launch flags + are unused.""" + k8, scale = reference_k_rows( + indexer_k, positions, cos_sin, ln_weight, ln_bias, eps, rope_dim, interleave, use_ue8m0 + ) + cache = k_cache.clone() + write_cache_rows(cache, k8, scale, slot_mapping_fp8, slot_mapping_scale, num_cache_tokens) + return k8, scale, cache + + def is_valid( + self, + indexer_k: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + eps: float, + rope_dim: int, + interleave: bool, + use_ue8m0: bool, + k_fp8: torch.Tensor, + k_scale: torch.Tensor, + k_cache: torch.Tensor, + slot_mapping_fp8: torch.Tensor, + slot_mapping_scale: torch.Tensor, + num_cache_tokens: int, + use_pdl: bool, + launch_dependents: bool, + prewait: bool, + ) -> None: + rows = indexer_k.shape[0] + # The kernel reads every operand through a unit inner stride and a fixed 128-wide row. + assert indexer_k.dim() == 2 and indexer_k.shape[1] == 128, ( + f"indexer_k must be [N, 128], got {tuple(indexer_k.shape)}" + ) + assert indexer_k.stride(1) == 1, "indexer_k needs a unit inner stride" + assert indexer_k.dtype in (torch.bfloat16, torch.float32), indexer_k.dtype + # The sin half starts rope_dim / 2 columns into the row: a wider or narrower table reads + # the wrong values without failing. + assert cos_sin.dtype == torch.float32 and cos_sin.dim() == 2, ( + "cos_sin must be a [max_pos, rope_dim] fp32 table" + ) + assert cos_sin.shape[1] == rope_dim and cos_sin.stride(1) == 1, ( + f"cos_sin rows must be [cos | sin] of rope_dim={rope_dim} contiguous values" + ) + assert rope_dim % 2 == 0 and 0 < rope_dim <= 128, f"rope_dim={rope_dim}" + assert positions.dim() == 1 and positions.numel() >= rows, ( + "positions must hold one position per row" + ) + if rows > 0: + assert int(positions[:rows].max()) < cos_sin.shape[0], ( + "a position past the cos/sin table reads past its end" + ) + assert ln_weight.numel() == 128 and ln_bias.numel() == 128, ( + "k_norm weight and bias hold 128 values" + ) + # The outputs are written as dense [N, 128] / [N, 1] rows. + assert k_fp8.shape == (rows, 128) and k_fp8.is_contiguous(), ( + f"k_fp8 must be a contiguous [N, 128] tensor, got {tuple(k_fp8.shape)}" + ) + assert k_fp8.dtype == torch.float8_e4m3fn, k_fp8.dtype + assert k_scale.shape == (rows, 1) and k_scale.is_contiguous(), ( + f"k_scale must be a contiguous [N, 1] tensor, got {tuple(k_scale.shape)}" + ) + assert k_scale.dtype == torch.float32, k_scale.dtype + assert k_cache.dim() == 4 and k_cache.element_size() == 1, ( + "k_cache must be the 4-D byte view of the indexer K cache" + ) + assert slot_mapping_fp8.dtype == torch.int64 and slot_mapping_scale.dtype == torch.int64, ( + "slot mappings are int64 flat byte offsets" + ) + assert 0 < num_cache_tokens <= rows, ( + f"num_cache_tokens={num_cache_tokens} must be in [1, {rows}]" + ) + assert slot_mapping_fp8.numel() >= num_cache_tokens, "too few fp8 slots" + assert slot_mapping_scale.numel() >= num_cache_tokens, "too few scale slots" + limit = k_cache.numel() + written = slot_mapping_fp8[:num_cache_tokens] + assert int(written.max()) + 128 <= limit, "an fp8 slot writes past the cache" + assert int(slot_mapping_scale[:num_cache_tokens].max()) + 4 <= limit, ( + "a scale slot writes past the cache" + ) + + +indexer_k_tail = GlmIndexerKTail() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/indexer_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/indexer_prologue.py new file mode 100644 index 000000000000..a8a9f032d90e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/indexer_prologue.py @@ -0,0 +1,272 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused indexer prologue: k_norm, RoPE, E4M3 Q and K, the scaled weights and the indexer K-cache +write in one Triton launch.""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._indexer_prologue import launch_indexer_prologue +from ._indexer_reference import reference_k_rows, reference_q_rows, write_cache_rows + +#: GLM-5.2's indexer: 32 heads of 128, the RoPE slice (first 64 values, interleaved pairs), the +#: k_norm eps, the UE8M0 scale format and softmax_scale * n_heads**-0.5. K and the weights are +#: column views of the fused-A output (kv_a_proj_with_mqa 2624 rows, wk 128, weights_proj 32). +_GLM = dict( + n_heads=32, + head_dim=128, + rope_dim=64, + interleave=True, + use_ue8m0=True, + eps=1e-6, + weight_scale=128**-0.5 * 32**-0.5, + tokens_per_block=64, + fused_a_width=2784, + k_col=2624, +) + + +def _spec(rows, first_position, **overrides) -> dict: + spec = dict( + _GLM, + rows=rows, + cache_rows=rows, + first_position=first_position, + k_layout="fused_a", + pool_layout="contiguous", + pages=16, + write_cache=True, + use_pdl=True, + launch_dependents=True, + prewait=True, + ) + spec.update(overrides) + return spec + + +class GlmIndexerPrologue(OpWrapper): + """The whole indexer prologue of a step whose Q comes from a separate wq_b GEMM. + + For every row ``i`` (the K part is ``indexer_k_tail``'s, including the cache write when + ``k_cache`` is given and ``num_cache_tokens`` > 0): + + - Q: RoPE on ``q[i, h, :rope_dim]`` at ``positions[i]``, rounded to bf16; per head + ``q_scale[i, h] = max(amax, 1e-12) / 448`` (rounded up to a power of two with + ``use_ue8m0``) and ``q_fp8[i, h] = e4m3(q * rcp(q_scale))``; + - weights: ``(weights[i, h] * q_scale[i, h]) * weight_scale`` in fp32. + + Returns ``(q_fp8 [N, H, D], k_fp8 [N, D], k_scale [N, 1], weights [N, H] fp32, + q_scale [N, H, 1])``, freshly allocated, the contract of the indexer's Q and K inputs to the + MQA logits. PDL flags as in ``indexer_k_tail``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode, 8 rows: wq_b output Q, fused-A K/weights, cache write, PDL + pre-wait", + spec=_spec(8, 529_981), + ), + Cell( + why="decode, 8 rows, plain launch", + spec=_spec(8, 1_048_560, use_pdl=False, launch_dependents=False, prewait=False), + ), + Cell( + why="no cache write (the caller scatters the K rows itself)", + spec=_spec(8, 300, write_cache=False, launch_dependents=False, prewait=False), + ), + Cell( + why="padded graph rows: 8 rows, the first 3 written to the cache", + spec=_spec(8, 60, cache_rows=3), + ), + Cell( + why="context chunk: 256 rows of one request over 4+ pages, contiguous K and weights", + spec=_spec( + 256, + 1_000, + k_layout="contiguous", + use_pdl=False, + launch_dependents=False, + prewait=False, + ), + ), + ) + + note = """ + Metadata consumed: none; every input is an argument (see ``indexer_k_tail`` for the K-side + operands). ``q`` is the indexer wq_b output viewed as [N, 32, 128] bf16 and ``weights`` the + weights_proj columns [2752, 2784) of the fused-A output (bf16). + + At the c1 decode cells the target takes the Q part from the wq_b epilogue and runs only + ``indexer_k_tail``; this entry serves the forwards that epilogue does not. + + Not certified: fp32 K or weights, NeoX RoPE, the plain amax / 448 scale, other head counts or + dims. + """ + + def __call__( + self, + indexer_k: torch.Tensor, + weights: torch.Tensor, + q: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + eps: float, + weight_scale: float, + rope_dim: int, + interleave: bool, + use_ue8m0: bool, + k_cache: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + num_cache_tokens: int, + use_pdl: bool, + launch_dependents: bool, + prewait: bool, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return launch_indexer_prologue( + indexer_k, + weights, + q, + positions, + cos_sin, + ln_weight, + ln_bias, + eps, + weight_scale, + rope_dim, + interleave, + use_ue8m0, + k_cache, + slot_mapping_fp8, + slot_mapping_scale, + num_cache_tokens, + use_pdl, + launch_dependents, + prewait, + ) + + def reference( + self, + indexer_k: torch.Tensor, + weights: torch.Tensor, + q: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + eps: float, + weight_scale: float, + rope_dim: int, + interleave: bool, + use_ue8m0: bool, + k_cache: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + num_cache_tokens: int, + use_pdl: bool, + launch_dependents: bool, + prewait: bool, + ) -> Tuple[torch.Tensor, ...]: + """``(q_fp8, k_fp8, k_scale, weights, q_scale, k_cache)``: the five outputs, then the cache + after the writes (a fresh tensor; None without ``k_cache``). The launch flags change when + the kernel runs, not what it computes.""" + q8, q_scale, w = reference_q_rows( + q, weights, positions, cos_sin, weight_scale, rope_dim, interleave, use_ue8m0 + ) + k8, k_scale = reference_k_rows( + indexer_k, positions, cos_sin, ln_weight, ln_bias, eps, rope_dim, interleave, use_ue8m0 + ) + cache = None + if k_cache is not None: + cache = k_cache.clone() + if num_cache_tokens > 0: + write_cache_rows( + cache, k8, k_scale, slot_mapping_fp8, slot_mapping_scale, num_cache_tokens + ) + return q8, k8, k_scale, w, q_scale, cache + + def is_valid( + self, + indexer_k: torch.Tensor, + weights: torch.Tensor, + q: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + eps: float, + weight_scale: float, + rope_dim: int, + interleave: bool, + use_ue8m0: bool, + k_cache: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + num_cache_tokens: int, + use_pdl: bool, + launch_dependents: bool, + prewait: bool, + ) -> None: + rows, n_heads, head_dim = q.shape + # Every operand is read through a unit inner stride and fixed 128-wide rows. + assert head_dim == 128 and indexer_k.shape == (rows, 128), ( + f"q must be [N, H, 128] and indexer_k [N, 128], got {tuple(q.shape)} and " + f"{tuple(indexer_k.shape)}" + ) + assert q.dtype == torch.bfloat16 and q.stride(2) == 1, "q must be bf16, unit inner stride" + assert indexer_k.stride(1) == 1 and weights.stride(1) == 1, ( + "indexer_k and weights need a unit inner stride" + ) + assert weights.shape == (rows, n_heads), ( + f"weights must be [N, H] = [{rows}, {n_heads}], got {tuple(weights.shape)}" + ) + assert indexer_k.dtype in (torch.bfloat16, torch.float32), indexer_k.dtype + assert weights.dtype in (torch.bfloat16, torch.float32), weights.dtype + assert cos_sin.dtype == torch.float32 and cos_sin.dim() == 2, ( + "cos_sin must be a [max_pos, rope_dim] fp32 table" + ) + # The sin half starts rope_dim / 2 columns into the row. + assert cos_sin.shape[1] == rope_dim and cos_sin.stride(1) == 1, ( + f"cos_sin rows must be [cos | sin] of rope_dim={rope_dim} contiguous values" + ) + assert rope_dim % 2 == 0 and 0 < rope_dim <= 128, f"rope_dim={rope_dim}" + assert positions.dim() == 1 and positions.numel() >= rows, ( + "positions must hold one position per row" + ) + if rows > 0: + assert int(positions[:rows].max()) < cos_sin.shape[0], ( + "a position past the cos/sin table reads past its end" + ) + assert ln_weight.numel() == 128 and ln_bias.numel() == 128, ( + "k_norm weight and bias hold 128 values" + ) + if k_cache is not None and num_cache_tokens > 0: + assert k_cache.dim() == 4 and k_cache.element_size() == 1, ( + "k_cache must be the 4-D byte view of the indexer K cache" + ) + assert slot_mapping_fp8 is not None and slot_mapping_scale is not None, ( + "a cache write needs both slot mappings" + ) + assert slot_mapping_fp8.dtype == torch.int64, "slot mappings are int64" + assert slot_mapping_scale.dtype == torch.int64, "slot mappings are int64" + assert num_cache_tokens <= rows, f"num_cache_tokens={num_cache_tokens} > {rows} rows" + assert slot_mapping_fp8.numel() >= num_cache_tokens, "too few fp8 slots" + assert slot_mapping_scale.numel() >= num_cache_tokens, "too few scale slots" + limit = k_cache.numel() + assert int(slot_mapping_fp8[:num_cache_tokens].max()) + 128 <= limit, ( + "an fp8 slot writes past the cache" + ) + assert int(slot_mapping_scale[:num_cache_tokens].max()) + 4 <= limit, ( + "a scale slot writes past the cache" + ) + + +indexer_prologue = GlmIndexerPrologue() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/paged_mqa_logits_schedule.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/paged_mqa_logits_schedule.py new file mode 100644 index 000000000000..c2eee988a5dc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/paged_mqa_logits_schedule.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-SM work partition of the indexer's paged MQA logits (DeepGEMM's paged metadata).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata + +from ._fp8_paged_mqa_logits import SPLIT_KV +from ._fp8_paged_mqa_logits import num_sms as _device_sms + +#: The indexer K cache's page size, which the schedule call takes as its block size. +_PAGE = 64 + + +def _spec(context_lens) -> dict: + return dict(context_lens=tuple(context_lens), block_kv=_PAGE) + + +class GlmPagedMqaLogitsSchedule(OpWrapper): + """``out`` [num_sms + 1, 2] int32: where each persistent CTA of the paged MQA logits starts. + + Request ``r`` with ``context_lens[r]`` keys is ``ceil(context_lens[r] / 256)`` units of work, + one per 256-key split (two 128-key compute blocks, one per math warp group), numbered request by + request. With ``total`` units, CTA ``s`` (0 <= s <= num_sms) starts at unit + ``w = s * (total // num_sms) + min(s, total % num_sms)``: ``out[s] = (r, w - first unit of r)`` + for the request ``r`` holding unit ``w``, and ``(num_requests, 0)`` past the last unit. CTA + ``s`` covers [out[s], out[s + 1]). + + Each request is one query atom (its whole verify window), so ``context_lens`` is passed as + [num_requests, 1]. Returns ``out``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="c1 key shard: one request, a full 33,152-key shard", spec=_spec([33_152])), + Cell(why="an empty shard: one request of one key", spec=_spec([1])), + Cell(why="one split per unit edge: 256 keys exactly", spec=_spec([256])), + Cell( + why="more splits than SMs: the whole 1M-key context on one rank", + spec=_spec([1_048_596]), + ), + Cell( + why="four requests, one shorter than a split", + spec=_spec([3_000, 70_000, 129, 40_000]), + ), + ) + + note = """ + Metadata consumed: none. The target passes the shard's context length (``hist_shard_meta``'s + ``shard_ctx``) or the generation requests' indexer KV lengths, the page size (64) as + ``block_kv``, the device's SM count (the logits kernel runs one CTA per SM) and a persistent + [num_sms + 1, 2] int32 ``out`` (captured CUDA graphs read it). The schedule kernel waits on + ``griddepcontrol.wait`` before it reads the lengths. + + ``block_kv`` must be 32, 64 or 128; it does not change the 256-key split on this + architecture. + """ + + def __call__( + self, context_lens: torch.Tensor, block_kv: int, num_sms: int, out: torch.Tensor + ) -> torch.Tensor: + out.copy_( + get_paged_mqa_logits_metadata(context_lens.view(-1, 1), block_kv, num_sms), + non_blocking=True, + ) + return out + + def reference( + self, context_lens: torch.Tensor, block_kv: int, num_sms: int, out: torch.Tensor + ) -> torch.Tensor: + """The partition, in Python integers; ``out`` only gives the result's device.""" + work = [-(-ctx // SPLIT_KV) for ctx in context_lens.tolist()] + prefix = [] + total = 0 + for units in work: + total += units + prefix.append(total) + quota, rem = divmod(total, num_sms) + rows = [] + for sm in range(num_sms + 1): + w = sm * quota + min(sm, rem) + request = next((r for r, end in enumerate(prefix) if end > w), len(work)) + if request < len(work): + rows.append((request, w - (prefix[request - 1] if request else 0))) + else: + rows.append((len(work), 0)) + return torch.tensor(rows, dtype=torch.int32, device=out.device) + + def is_valid( + self, context_lens: torch.Tensor, block_kv: int, num_sms: int, out: torch.Tensor + ) -> None: + # Each element is read as one request's length. + assert context_lens.dim() == 1, "context_lens must be [num_requests]" + # The schedule is right for any SM count, but the logits kernel launches one CTA per SM + # and reads entries [s] and [s + 1]: another count leaves keys unscored or reads past the + # schedule. + device = context_lens.device.index if context_lens.device.index is not None else 0 + assert num_sms == _device_sms(device), ( + f"num_sms={num_sms} differs from the device's SM count, the logits kernel's CTA count" + ) + # A wider int type would be cast silently by the copy. + assert out.dtype == torch.int32, f"out must be int32, got {out.dtype}" + + +paged_mqa_logits_schedule = GlmPagedMqaLogitsSchedule() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/stash_last_accepted_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/stash_last_accepted_rows.py new file mode 100644 index 000000000000..02b6d14a1300 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/stash_last_accepted_rows.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MTP index share: draft step 0's TopK row of each request's last accepted token, in one launch.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._mtp_topk_stash import launch_stash_last_accepted_rows + + +def _spec(requests, accepted, **overrides) -> dict: + spec = dict(requests=requests, next_n=8, n_cols=2048, accepted=tuple(accepted), dst="rows") + spec.update(overrides) + return spec + + +class GlmStashLastAcceptedRows(OpWrapper): + """``dst[r] = gen_topk[r * next_n + clamp(num_accepted[r] - 1, 0, next_n - 1)]`` per request. + + Draft step 0 of the MTP draft loop computes the TopK of the verify step's ``next_n`` rows; the + draft steps after it reuse, per request, the row of its last accepted token. ``num_accepted`` + is the verify step's accepted-token count per request (1 to next_n); out-of-range counts clamp. + Returns ``dst``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="c1: one request, 8 TopK rows, 3 tokens accepted", spec=_spec(1, [3])), + Cell(why="c1: every draft token accepted (the last row)", spec=_spec(1, [8])), + Cell(why="counts past the window clamp: 0 and 9", spec=_spec(2, [0, 9])), + Cell( + why="four requests into rows of a wider shared buffer", + spec=_spec(4, [1, 5, 8, 2], dst="wide"), + ), + ) + + note = """ + Metadata consumed: none. The target passes the draft step-0 TopK output + ([num_requests * 8, 2048] int32), the verify step's accepted counts per request (int32 on the + device) and the rows [0, num_requests) of the persistent shared-TopK buffer. A plain launch + (no PDL): at c1 it runs on its own stream behind an event recorded after the TopK, and the + layer joins it at the end of its attention. + """ + + def __call__( + self, gen_topk: torch.Tensor, num_accepted: torch.Tensor, dst: torch.Tensor, next_n: int + ) -> torch.Tensor: + launch_stash_last_accepted_rows(gen_topk, num_accepted, dst, next_n) + return dst + + def reference( + self, gen_topk: torch.Tensor, num_accepted: torch.Tensor, dst: torch.Tensor, next_n: int + ) -> torch.Tensor: + """The selected rows as a fresh tensor.""" + requests = dst.shape[0] + offsets = (num_accepted.long() - 1).clamp(0, next_n - 1) + index = torch.arange(requests, device=gen_topk.device) * next_n + offsets + return gen_topk[index, : dst.shape[1]].clone() + + def is_valid( + self, gen_topk: torch.Tensor, num_accepted: torch.Tensor, dst: torch.Tensor, next_n: int + ) -> None: + requests, n_cols = dst.shape + # Rows are addressed as request * next_n + offset with unit column strides. + assert gen_topk.shape == (requests * next_n, n_cols), ( + f"gen_topk must be [{requests} * {next_n}, {n_cols}], got {tuple(gen_topk.shape)}" + ) + assert gen_topk.dtype == dst.dtype, (gen_topk.dtype, dst.dtype) + assert gen_topk.stride(1) == 1 and dst.stride(1) == 1, "unit column strides" + assert num_accepted.is_contiguous() and num_accepted.numel() >= requests, ( + "num_accepted must hold one contiguous count per request" + ) + + +stash_last_accepted_rows = GlmStashLastAcceptedRows() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/embedding/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/embedding/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/embedding/fused_embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/embedding/fused_embedding.py new file mode 100644 index 000000000000..746926a8b680 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/embedding/fused_embedding.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Token embedding lookup from a replicated BF16 table in one kernel, via the trtllm +glm_fused_embedding op.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: GLM-5.2's vocabulary and hidden size: the table every rank holds whole. +VOCAB = 154880 +HIDDEN = 6144 + + +class GlmFusedEmbedding(OpWrapper): + """`out[..., :] = weight[input_ids[...], :]`, and a zero row for an id outside + `[0, weight.shape[0])`. + + The guard is the Embedding module's for an unsharded table: the -1 that + pads rejected draft tokens, or any other out-of-range id, reads as a zero + row instead of out of bounds. Rows are bit copies of the table. Returns a + new tensor of shape `input_ids.shape + (hidden,)` in the table's dtype; no + input is modified. + + `weight` is the whole table, [vocab, hidden], the same on every rank: the + lookup is local and needs no collective. A vocab-parallel shard is not a + valid `weight` (see `note`). + + One kernel on the current stream: one thread per 16-byte vector, a grid of + tokens x hidden / 1024 CTAs of 128 threads (8 x 6 at the c1 verify + forward). Under programmatic dependent launch it waits for its producer + before reading the ids, then releases its dependents at once and writes + `out`, so the kernel after it must wait before reading `out`. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The target looks up the verify forward's ids (one request's last accepted + # token and its 7 drafts at c1) in the replicated table; the engine's ids + # are int32. Prefill forwards (gen-server warm-up, PrefillTarget) look up a + # whole chunk. + CELLS: tuple[Cell, ...] = ( + Cell( + why="c1 verify forward: 8 int32 ids into the replicated 154880 x 6144 table", + spec=dict(tokens=8, ids_dtype=torch.int32, ids="random"), + ), + Cell( + why="decode without speculation: 1 id, a grid of 1 x 6 CTAs", + spec=dict(tokens=1, ids_dtype=torch.int32, ids="random"), + ), + Cell( + why="rejected-draft padding: -1 and ids >= vocab read as zero rows; rows 0, vocab-1", + spec=dict(tokens=8, ids_dtype=torch.int32, ids="edges"), + ), + Cell( + why="int64 ids: the second id type the kernel is instantiated for", + spec=dict(tokens=8, ids_dtype=torch.int64, ids="edges"), + ), + Cell( + why="prefill chunk: 8192 ids, a grid of 8192 x 6 CTAs", + spec=dict(tokens=8192, ids_dtype=torch.int32, ids="random"), + ), + ) + + note = """ + `weight`: BF16, contiguous, 2-D, 16-byte aligned, hidden a multiple of 8; + ids int32 or int64 on the table's device, any shape (made contiguous). The + op raises on anything else. + + A vocab-parallel shard passed as `weight` is not detectable from the + tensors and answers wrongly: the op indexes it with the global ids (rows + of the wrong tokens, zeros past the shard) instead of masking to the + shard's range and summing over ranks. The target keeps the table whole. + + Bitwise equal to the Embedding module's masked F.embedding chain for an + unsharded table. + """ + + def __call__(self, input_ids: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + return torch.ops.trtllm.glm_fused_embedding(input_ids, weight) + + def reference(self, input_ids: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """The table rows of the in-range ids, zero rows elsewhere: a gather, + nothing accumulated or rounded.""" + valid = (input_ids >= 0) & (input_ids < weight.shape[0]) + rows = torch.where(valid, input_ids, torch.zeros_like(input_ids)).long() + return torch.nn.functional.embedding(rows, weight).masked_fill(~valid.unsqueeze(-1), 0) + + def is_valid(self, input_ids: torch.Tensor, weight: torch.Tensor) -> None: + """Nothing to refuse: every input the op accepts it computes as stated. + + The op raises on every dtype, layout, alignment and device it does not + serve. What is left is whether `weight` is the whole table (see + `note`), which no check on the tensors can see. + """ + return None + + +fused_embedding = GlmFusedEmbedding() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk.py new file mode 100644 index 000000000000..8c592c735b35 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk.py @@ -0,0 +1,1890 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Adapted from FlashInfer (https://github.com/flashinfer-ai/flashinfer) under the Apache-2.0 +# license. Original file: flashinfer/gemm/kernels/dense_bf16_gemm_sm100_splitk.py +# Original copyright: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +"""SM10x low-M BF16/FP16 GEMM with an in-kernel cluster split-K reduction (CuTe DSL). + +Each cluster rank accumulates an exact K slice in FP32. Peers publish partials +to rank 0 through DSMEM; rank 0 reduces, casts, and stores once. The public +``A[M, K] @ B[K, N]`` problem is swapped internally, so tile dimensions below +use kernel coordinates: kernel-M carries public N and kernel-N carries public M. + +This target's own copy of the kernel. Beside the plain GEMM it carries what the +target's decode step needs and FlashInfer's kernel does not have: + +- ``splitk_default_tactic``: the source default tactic, refined to a cluster of + eight K slices where that fills otherwise idle SMs; +- ``LamportPublish`` (``run_splitk_publish``): the output stored straight into + this rank's slot of an MNNVL one-shot all-reduce workspace; +- ``HnormB`` and an FP32 output (``run_splitk_hidden_half``): B produced in the + kernel as FlashInfer's RMSNorm of one activation row, the result kept in FP32. + +The launchers validate their operands in Python (``ValueError``) and compile one +kernel per specialization on first use. +""" + +from __future__ import annotations + +import dataclasses +import functools +import struct + +import cuda.bindings.driver as _cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import torch as _torch +from cutlass import Float32, Int32, Int64 +from cutlass._mlir.dialects import llvm +from cutlass.cute import experimental as cute_ext +from cutlass.cute.nvgpu import tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +#: Per-CTA SMEM capacity reported by CuTeDSL on SM100/SM103. +_SMEM_CAPACITY_BYTES = utils.get_smem_capacity_in_bytes("sm_100") + +#: K extent of one CTA tile. +_CTA_K = 128 + +#: Kernel-M tiles; 64 increases CTA count for low-M decode shapes. +_SUPPORTED_MMA_M = (64, 128) + +#: Kernel-N carries public M, which is limited to 32. +_SUPPORTED_MMA_N = (8, 16, 32) + +#: Physical cluster-K sizes; split 1 compiles out the DSMEM path. +_SUPPORTED_SPLIT_K = (1, 2, 3, 4) + +#: Largest public M this low-M policy serves. +MAX_M = 32 + +#: Bytes per FP32 partial exchanged through DSMEM. +_FP32_BYTES = 4 + +#: DSMEM mailbox base alignment, in bytes. +_MAILBOX_ALIGN_BYTES = 128 + +#: Size and alignment of one mbarrier, in bytes. +_MBARRIER_BYTES = 8 + +#: Size and alignment of the TMEM base pointer slot. +_TMEM_POINTER_BYTES = 4 + +#: Bytes per BF16/FP16 element. +_AB_ELEMENT_BYTES = 2 + +#: Alignment of the A/B shared-memory buffers. +_AB_BUFFER_ALIGN_BYTES = 1024 + +#: A/B pipeline stage bounds. One stage is useful only for a single K tile. +_MIN_AB_STAGES = 1 +_DEFAULT_MIN_AB_STAGES = 2 +_MAX_AB_STAGES = 12 + + +@dataclasses.dataclass(frozen=True) +class SplitKTactic: + """One specialization; mma_m carries public N and mma_n carries public M.""" + + mma_m: int + mma_n: int + split_k: int + ab_stages: int + + +def _align_up(value: int, alignment: int) -> int: + return ((value + alignment - 1) // alignment) * alignment + + +def _smem_bytes( + tactic: SplitKTactic, + ab_stages: int, +) -> int: + """Mirror the device allocator's shared-memory layout.""" + cursor = ( + _align_up( + tactic.mma_m * _CTA_K * _AB_ELEMENT_BYTES * ab_stages, + _AB_BUFFER_ALIGN_BYTES, + ) + + tactic.mma_n * _CTA_K * _AB_ELEMENT_BYTES * ab_stages + ) + + cursor = _align_up(cursor, _MBARRIER_BYTES) + cursor += 2 * ab_stages * _MBARRIER_BYTES + cursor += 3 * _MBARRIER_BYTES + cursor = _align_up(cursor, _TMEM_POINTER_BYTES) + cursor += _TMEM_POINTER_BYTES + + if tactic.split_k == 1: + return cursor + + return ( + _align_up( + _align_up(cursor, _MAILBOX_ALIGN_BYTES) + + (tactic.split_k - 1) * tactic.mma_m * tactic.mma_n * _FP32_BYTES, + _MBARRIER_BYTES, + ) + + _MBARRIER_BYTES + ) + + +#: Headroom over the mirrored layout when a launch reserves only the tactic's footprint. +_LAUNCH_SMEM_HEADROOM_BYTES = 2048 + +#: Alignment of the Lamport-publish staging tile. +_PUBLISH_ALIGN_BYTES = 16 + +#: Named barrier of the four epilogue warps in the Lamport-publish epilogue (0 is the CTA-wide one). +_PUBLISH_BARRIER_ID = 1 + +#: BF16 elements per 16-byte Lamport-publish store. +_PUBLISH_VEC_ELEMENTS = 8 + + +def _publish_smem_bytes(tactic: SplitKTactic, ab_stages: int) -> int: + """Shared memory of the Lamport-publish kernel: the base layout, then the FP32 output tile.""" + return ( + _align_up(_smem_bytes(tactic, ab_stages), _PUBLISH_ALIGN_BYTES) + + tactic.mma_m * tactic.mma_n * _FP32_BYTES + ) + + +@dataclasses.dataclass(frozen=True) +class LamportPublish: + """Publish the GEMM output into the current buffer of an MNNVL one-shot all-reduce workspace + instead of writing C. + + The owner CTA of each output tile stores its BF16 values, -0.0 written as +0.0, through the + workspace's multicast address at [token][rank][column] of the buffer the workspace flags mark + current: the shard store of the one-shot kernel, which then only polls. Every field is a + compile-time constant of the kernel. + + ``trigger_after_wait``: the PDL dependents (the all-reduce) are released by the epilogue warps + right after the grid dependency wait instead of by the weight-load warp once its loads are + issued, so that the all-reduce's reads before its own wait (the flags, the residual) follow + the completion of this GEMM's primary. + """ + + world_size: int + rank: int + trigger_after_wait: bool = True + + +#: Row length whose FlashInfer RMSNorm reduction ``HnormB`` reproduces. +HNORM_B_HIDDEN = 6144 + +#: BF16 elements per 16-byte vector of the HnormB producer. +_HNORM_B_VEC_ELEMENTS = 8 + +#: Lanes of a warp. +_WARP_LANES = 32 + + +@dataclasses.dataclass(frozen=True) +class HnormB: + """Produce B in the kernel instead of loading it: B is FlashInfer's RMSNorm + (``flashinfer.norm.rmsnorm``, its CuTe DSL kernel at H = 6144) of the one activation row + passed as B, with the weight passed as the bias operand, rounded to BF16. Warp 1 computes the + norm with that kernel's arithmetic and summation order and writes this CTA's K slice into the + B stages, so the MMA reads what the TMA would have loaded from the norm's output and the GEMM's + result is bit-identical to the GEMM of that output. Every field is a compile-time constant. + """ + + hidden: int + eps: float + + +def hnorm_b_supported(tactic: SplitKTactic, hidden: int) -> bool: + """Whether ``HnormB`` serves ``tactic`` for a row of ``hidden``: the reproduced width, 8-row B + tiles (one K_SW128 swizzle atom row group per 64 columns), every K tile of a CTA in its own + stage, and a CTA slice of whole 32-vector groups (one warp, one vector per lane per group).""" + if hidden != HNORM_B_HIDDEN or tactic.mma_n != 8: + return False + if hidden % (_CTA_K * tactic.split_k): + return False + k_tiles = hidden // _CTA_K // tactic.split_k + return ( + k_tiles <= tactic.ab_stages + and (k_tiles * _CTA_K // _HNORM_B_VEC_ELEMENTS) % _WARP_LANES == 0 + ) + + +def tactic_launch_smem_bytes(tactic: SplitKTactic) -> int: + """Dynamic shared memory for a launch that reserves only what ``tactic`` uses. + + The default launch reserves the full capacity, which keeps one CTA per SM. This size + lets other CTAs share the SM: two CTAs of the c1 q_b tactic (128, 8, 8, 2) fit on one. + """ + return min( + _align_up(_smem_bytes(tactic, tactic.ab_stages), 1024) + _LAUNCH_SMEM_HEADROOM_BYTES, + _SMEM_CAPACITY_BYTES, + ) + + +def _max_ab_stages_for( + tactic: SplitKTactic, + smem_capacity: int, +) -> int: + return next( + ( + stages + for stages in range(_MAX_AB_STAGES, -1, -1) + if _smem_bytes(tactic, stages) <= smem_capacity + ), + 0, + ) + + +def validate_tactic( + tactic: SplitKTactic, + m: int, + n: int, + k: int, + *, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> None: + """Reject a tactic that cannot serve ``(m, n, k)``.""" + if tactic.mma_m not in _SUPPORTED_MMA_M: + raise ValueError(f"unsupported mma_m={tactic.mma_m}") + if tactic.mma_n not in _SUPPORTED_MMA_N: + raise ValueError(f"unsupported mma_n={tactic.mma_n}") + # Split 8 (splitk_default_tactic) is served beside the cataloged cluster sizes. + if tactic.split_k not in _SUPPORTED_SPLIT_K and tactic.split_k != 8: + raise ValueError(f"unsupported split_k={tactic.split_k}") + if not _MIN_AB_STAGES <= tactic.ab_stages <= _MAX_AB_STAGES: + raise ValueError( + f"ab_stages must be in [{_MIN_AB_STAGES}, {_MAX_AB_STAGES}], got {tactic.ab_stages}" + ) + if not 1 <= m <= MAX_M: + raise ValueError(f"this low-M policy requires 1 <= M <= {MAX_M}, got {m}") + if n <= 0: + raise ValueError(f"N must be positive, got {n}") + if k <= 0 or k % _CTA_K or (k // _CTA_K) % tactic.split_k: + raise ValueError( + f"K={k} with CTA_K={_CTA_K} does not divide evenly across split_k={tactic.split_k}" + ) + smem_bytes = _smem_bytes(tactic, tactic.ab_stages) + if smem_bytes > smem_capacity: + raise ValueError( + f"tactic {tactic} needs {smem_bytes} B of shared memory but only " + f"{smem_capacity} B are available; max ab_stages is " + f"{_max_ab_stages_for(tactic, smem_capacity)}" + ) + + +@functools.cache +def device_limits(device_index: int) -> tuple[int, int]: + """(SM count, opt-in shared memory per block) of a CUDA device, the resources + ``splitk_default_tactic`` plans for.""" + properties = _torch.cuda.get_device_properties(device_index) + return properties.multi_processor_count, properties.shared_memory_per_block_optin + + +def default_tactic(m: int, n: int, k: int) -> SplitKTactic: + """Choose the default occupancy-oriented tactic.""" + if n <= 512: + mma_m = 64 + mma_n = 16 if n == 512 and m > 24 else 8 + requested_split = 4 + else: + mma_n = 8 if m <= 8 else 16 if m <= 16 else 32 + if n <= 3072: + mma_m = 128 if m <= 16 else 64 + requested_split = 4 if m <= 16 else 2 + elif n < 8192: + mma_m = 64 + requested_split = 2 + else: + mma_m = 128 if k <= 1024 and m <= 24 else 64 + requested_split = 1 + + if k <= 4 * _CTA_K: + requested_split = 1 + split_k = next( + split_k + for split_k in reversed(_SUPPORTED_SPLIT_K) + if split_k <= requested_split and (k // _CTA_K) % split_k == 0 + ) + tactic = SplitKTactic(mma_m, mma_n, split_k, _MIN_AB_STAGES) + max_stages = _max_ab_stages_for(tactic, _SMEM_CAPACITY_BYTES) + tactic = dataclasses.replace( + tactic, + ab_stages=( + _MIN_AB_STAGES + if k == _CTA_K + else _DEFAULT_MIN_AB_STAGES + if k <= 2 * _CTA_K and m > 8 + else min(max_stages, 6) + if k <= 4 * _CTA_K + else max_stages + ), + ) + validate_tactic(tactic, m, n, k) + return tactic + + +def splitk_default_tactic( + m: int, n: int, k: int, *, sm_count: int, smem_capacity: int +) -> SplitKTactic: + """The source default tactic, refined to split 8 where split 8 fills otherwise idle SMs. + + The refinement needs exact K slices of at least two K tiles per peer and at most one wave + of eight-CTA clusters; other shapes keep the source default. An occupancy heuristic, not a + measured winner for every shape. The launch reserves the full SM100 shared-memory capacity, + even when the selected pipeline uses fewer bytes. + """ + if sm_count <= 0 or smem_capacity < _SMEM_CAPACITY_BYTES: + raise ValueError("device cannot serve the split-K launch resources") + tactic = default_tactic(m, n, k) + if tactic.split_k > sm_count: + raise ValueError("device has too few SMs for the split-K cluster") + output_tiles = ((n + tactic.mma_m - 1) // tactic.mma_m) * ( + (m + tactic.mma_n - 1) // tactic.mma_n + ) + # Require exact K slices, at least two tiles per peer, and at most one + # wave of eight-CTA clusters. Wide outputs retain their existing tactic. + if ( + k % (8 * _CTA_K) + or k // (8 * _CTA_K) < _DEFAULT_MIN_AB_STAGES + or output_tiles * tactic.split_k >= sm_count + or output_tiles * 8 > sm_count + ): + return tactic + candidate = dataclasses.replace(tactic, split_k=8) + stages = min( + tactic.ab_stages, + k // (8 * _CTA_K), + _max_ab_stages_for(candidate, _SMEM_CAPACITY_BYTES), + ) + if stages < _DEFAULT_MIN_AB_STAGES: + return tactic + candidate = dataclasses.replace(candidate, ab_stages=stages) + validate_tactic(candidate, m, n, k) + return candidate + + +__all__ = [ + "MAX_M", + "HnormB", + "LamportPublish", + "SplitKTactic", + "default_tactic", + "device_limits", + "hidden_half_norm_supported", + "hnorm_b_supported", + "operands_overlap", + "publish_supported", + "publish_tactic", + "run_splitk_dense", + "run_splitk_hidden_half", + "run_splitk_publish", + "splitk_default_tactic", + "tactic_launch_smem_bytes", + "validate_tactic", +] + + +@dsl_user_op +def _map_shared_rank( + smem_ptr: "cute.Pointer", + peer_cta_rank_in_cluster: "Int32", + *, + loc=None, + ip=None, +) -> "Int32": + """Map an SMEM pointer into a peer CTA's address space.""" + return Int32( + llvm.inline_asm( + T.i32(), + [ + smem_ptr.toint(loc=loc, ip=ip).ir_value(), + peer_cta_rank_in_cluster.ir_value(), + ], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _store_shared_remote_v4( + value0, + value1, + value2, + value3, + smem_ptr: "cute.Pointer", + mbar_ptr: "cute.Pointer", + peer_cta_rank_in_cluster: "Int32", + *, + loc=None, + ip=None, +) -> None: + """Publish four FP32 partials into a peer's SMEM, crediting 16 bytes.""" + llvm.inline_asm( + None, + [ + _map_shared_rank(smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), + value0.bitcast(Int32).ir_value(loc=loc, ip=ip), + value1.bitcast(Int32).ir_value(loc=loc, ip=ip), + value2.bitcast(Int32).ir_value(loc=loc, ip=ip), + value3.bitcast(Int32).ir_value(loc=loc, ip=ip), + _map_shared_rank(mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), + ], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 " + "[$0], {$1, $2, $3, $4}, [$5];", + "r,r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _lamport_current_buffer( + flags_address: "Int64", + buffers_address: "Int64", + *, + loc=None, + ip=None, +) -> "Int64": + """The address of the workspace's current Lamport buffer: buffers + flags[0] (current index) + * flags[2] (bytes per buffer), the flags read with one volatile 16-byte load (the previous + all-reduce rotated them).""" + return Int64( + llvm.inline_asm( + T.i64(), + [flags_address.ir_value(loc=loc, ip=ip), buffers_address.ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 lfcur, lfdirty, lfbytes, lfstages;\n\t.reg .b64 lfoffset;\n\t" + "ld.volatile.global.v4.u32 {lfcur, lfdirty, lfbytes, lfstages}, [$1];\n\t" + "mul.wide.u32 lfoffset, lfcur, lfbytes;\n\tadd.s64 $0, $2, lfoffset;\n\t}", + "=l,l,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _cvt_bf16x2_lamport(lo: "Float32", hi: "Float32", *, loc=None, ip=None) -> "Int32": + """Round two FP32 values to BF16 (nearest even), ``lo`` in the low half, and write a -0.0 half + (0x8000, the idle value of a Lamport buffer) as +0.0, as the one-shot all-reduce does before it + publishes its shard.""" + return Int32( + llvm.inline_asm( + T.i32(), + [lo.ir_value(loc=loc, ip=ip), hi.ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 lpword, lplo, lphi;\n\t.reg .pred lpzlo, lpzhi;\n\t" + "cvt.rn.bf16x2.f32 lpword, $2, $1;\n\t" + "and.b32 lplo, lpword, 0xFFFF;\n\tshr.u32 lphi, lpword, 16;\n\t" + "setp.eq.u32 lpzlo, lplo, 0x8000;\n\tsetp.eq.u32 lpzhi, lphi, 0x8000;\n\t" + "@lpzlo and.b32 lpword, lpword, 0xFFFF0000;\n\t" + "@lpzhi and.b32 lpword, lpword, 0xFFFF;\n\t" + "mov.b32 $0, lpword;\n\t}", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _store_global_v4( + address: "Int64", v0: "Int32", v1: "Int32", v2: "Int32", v3: "Int32", *, loc=None, ip=None +) -> None: + """One 16-byte global store (also through a multicast address).""" + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip)] + [v.ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "st.global.v4.b32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _ld_global_cg_v4(address: "Int64", *, loc=None, ip=None): + """One 16-byte load through L2 (``ld.global.cg``): four words.""" + res = llvm.inline_asm( + llvm.StructType.get_literal([T.i32(), T.i32(), T.i32(), T.i32()]), + [address.ir_value(loc=loc, ip=ip)], + "ld.global.cg.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + return [Int32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(4)] + + +@dsl_user_op +def _st_shared_v4( + address: "Int32", v0: "Int32", v1: "Int32", v2: "Int32", v3: "Int32", *, loc=None, ip=None +) -> None: + """One 16-byte shared-memory store.""" + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip)] + [v.ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "st.shared.v4.b32 [$0], {$1, $2, $3, $4};", + "r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _fence_proxy_async_shared(*, loc=None, ip=None) -> None: + """Order this thread's shared-memory stores before later async-proxy (tensor core) reads.""" + llvm.inline_asm( + None, + [], + "fence.proxy.async.shared::cta;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _add_bf16x2_squares(acc: "Float32", word: "Int32", *, loc=None, ip=None) -> "Float32": + """``(acc + lo * lo) + hi * hi`` for the two BF16 halves of ``word`` (low half first), every + product and sum rounded to nearest, no FMA: FlashInfer RMSNorm's per-thread accumulation.""" + return Float32( + llvm.inline_asm( + T.f32(), + [acc.ir_value(loc=loc, ip=ip), word.ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 sqlo, sqhi;\n\t.reg .f32 sqa, sqb, sqp, sqs;\n\t" + "shl.b32 sqlo, $2, 16;\n\tand.b32 sqhi, $2, 0xFFFF0000;\n\t" + "mov.b32 sqa, sqlo;\n\tmov.b32 sqb, sqhi;\n\t" + "mul.rn.f32 sqp, sqa, sqa;\n\tadd.rn.f32 sqs, $1, sqp;\n\t" + "mul.rn.f32 sqp, sqb, sqb;\n\tadd.rn.f32 $0, sqs, sqp;\n\t}", + "=f,f,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _add_rn(a: "Float32", b: "Float32", *, loc=None, ip=None) -> "Float32": + """``a + b`` rounded to nearest (never contracted).""" + return Float32( + llvm.inline_asm( + T.f32(), + [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +def _f32_hex(value: float) -> str: + """PTX literal of the FP32 rounding of ``value``.""" + return "0f%08X" % struct.unpack(" "Float32": + """FlashInfer RMSNorm's ``rsqrt.approx.ftz(total / hidden + eps)``, the division and the sum + rounded to nearest (IEEE).""" + return Float32( + llvm.inline_asm( + T.f32(), + [total.ir_value(loc=loc, ip=ip)], + "{\n\t.reg .f32 rsm, rst;\n\t" + f"div.rn.f32 rsm, $1, {_f32_hex(float(hidden))};\n\t" + f"add.rn.f32 rst, rsm, {_f32_hex(eps)};\n\t" + "rsqrt.approx.ftz.f32 $0, rst;\n\t}", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _rmsnorm_bf16x2( + x_word: "Int32", w_word: "Int32", rstd: "Float32", *, loc=None, ip=None +) -> "Int32": + """FlashInfer RMSNorm's output for two BF16 elements (low half first): ``(x * rstd) * (w + 0)`` + rounded to nearest at every step, then to BF16 (nearest even).""" + return Int32( + llvm.inline_asm( + T.i32(), + [ + x_word.ir_value(loc=loc, ip=ip), + w_word.ir_value(loc=loc, ip=ip), + rstd.ir_value(loc=loc, ip=ip), + ], + "{\n\t.reg .b32 rnxl, rnxh, rnwl, rnwh;\n\t" + ".reg .f32 rfxl, rfxh, rfwl, rfwh, ryl, ryh;\n\t" + "shl.b32 rnxl, $1, 16;\n\tand.b32 rnxh, $1, 0xFFFF0000;\n\t" + "shl.b32 rnwl, $2, 16;\n\tand.b32 rnwh, $2, 0xFFFF0000;\n\t" + "mov.b32 rfxl, rnxl;\n\tmov.b32 rfxh, rnxh;\n\t" + "mov.b32 rfwl, rnwl;\n\tmov.b32 rfwh, rnwh;\n\t" + "add.rn.f32 rfwl, rfwl, 0f00000000;\n\tadd.rn.f32 rfwh, rfwh, 0f00000000;\n\t" + "mul.rn.f32 ryl, rfxl, $3;\n\tmul.rn.f32 ryl, ryl, rfwl;\n\t" + "mul.rn.f32 ryh, rfxh, $3;\n\tmul.rn.f32 ryh, ryh, rfwh;\n\t" + "cvt.rn.bf16x2.f32 $0, ryh, ryl;\n\t}", + "=r,r,r,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +#: Rank that gathers partials and stores the output. +OWNER_RANK = 0 + + +class SplitKDenseGemmKernel: + """Standalone BF16/FP16 GEMM with a cluster-local split-K reduction.""" + + def __init__( + self, + *, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + reserve_full_smem: bool = True, + publish: LamportPublish | None = None, + hnorm_b: HnormB | None = None, + ) -> None: + self.acc_dtype = cutlass.Float32 + self.cta_m = tactic.mma_m + self.cta_n = tactic.mma_n + self.cta_k = _CTA_K + self.num_ab_stage = tactic.ab_stages + self.split_k = tactic.split_k + self.use_pdl = use_pdl + self.has_bias = has_bias + # Shared memory requested at launch: the full capacity (one CTA per SM) or only the + # tactic's footprint. The kernel itself is the same either way. + self.reserve_full_smem = reserve_full_smem + self.launch_smem_bytes = tactic_launch_smem_bytes(tactic) + # Lamport publish (run_splitk_publish): the owner stages the reduced FP32 tile in shared + # memory and stores it into the all-reduce workspace instead of writing C. + self.publish = publish + # HnormB: warp 1 computes B (FlashInfer's RMSNorm of the activation row) instead of + # loading it; the norm weight takes the bias operand's place. + self.hnorm_b = hnorm_b + if hnorm_b is not None and not hnorm_b_supported(tactic, hnorm_b.hidden): + raise ValueError( + f"HnormB needs hidden {HNORM_B_HIDDEN} and a tactic of hnorm_b_supported: " + f"{tactic}, hidden {hnorm_b.hidden}" + ) + if hnorm_b is not None and (has_bias or publish is not None): + raise ValueError("HnormB excludes a bias and the Lamport publish") + if publish is not None: + if has_bias or self.cta_m % _PUBLISH_VEC_ELEMENTS: + raise ValueError(f"Lamport publish needs no bias and 8-column tiles: {tactic}") + self.launch_smem_bytes = min( + _align_up(_publish_smem_bytes(tactic, tactic.ab_stages), 1024) + + _LAUNCH_SMEM_HEADROOM_BYTES, + _SMEM_CAPACITY_BYTES, + ) + + self.threads_per_cta = 256 + self.epilog_threads = 128 + self.mma_tiler_mn = (tactic.mma_m, tactic.mma_n) + self.cta_group = tcgen05.CtaGroup.ONE + self.tma_op = cute_ext.OperationTypeEnum.SM90_TMA_LOAD + self.cluster_shape = (1, tactic.split_k, 1) + + values_per_thread = (tactic.mma_m * tactic.mma_n) // self.epilog_threads + if values_per_thread % 4: + raise ValueError( + f"CTA tile ({tactic.mma_m}, {tactic.mma_n}) gives " + f"{values_per_thread} " + "values per epilogue thread; remote stores require a multiple of 4" + ) + self.mailbox_elements = (tactic.split_k - 1) * self.epilog_threads * values_per_thread + self.expected_transaction_bytes = self.mailbox_elements * _FP32_BYTES + + @cute.experimental.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + lamport: cute.Tensor, + flags: cute.Tensor, + stream: _cuda.CUstream, + ): + # Grid-y packs output-N tile and cluster rank. + self.kernel(a, b, c, bias, lamport, flags).launch( + grid=( + cute.ceil_div(c.layout.shape[0], self.cta_m), + cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k, + c.layout.shape[2], + ), + block=(self.threads_per_cta, 1, 1), + cluster=self.cluster_shape, + smem=cute.Int64( + utils.get_smem_capacity_in_bytes("sm_100") + if self.reserve_full_smem + else self.launch_smem_bytes + ), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.experimental.kernel + def kernel( + self, + mA: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L), K-major + mB: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L), K-major + mC: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L), M-major + mBias: cute.Tensor, # Broadcast bias; dead when has_bias=False + mLamport: cute.Tensor, # BF16 all-reduce workspace, multicast address; dead without publish + mFlags: cute.Tensor, # The workspace's buffer flags; dead without publish + ): + """Allocate storage and dispatch the specialized warps.""" + stages = self.num_ab_stage + + ab_dtype = mA.element_type + tiled_mma = sm100_utils.make_trivial_tiled_mma( + ab_dtype, + ab_dtype, + utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + + mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k) + block_idx = cute.arch.block_idx() + bidx = block_idx[0] + split_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + n_idx = block_idx[1] // self.split_k + l_idx = block_idx[2] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + sA = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + sB = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + + acc_layout = cute_ext.make_tmem_layout_acc(tiled_mma, self.mma_tiler_mn, acc_stage=1) + c_tiler_mn = (self.cta_m, self.cta_n) + + bar_full = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_empty = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_mma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tmem_alloc = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + tmem_base_ptr = cute_ext.allocate( + cutlass.Int32, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_TMEM_POINTER_BYTES, + ).iterator + + if cutlass.const_expr(self.split_k > 1): + mailbox = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.mailbox_elements), + alignment=_MAILBOX_ALIGN_BYTES, + ) + bar_reduce = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + else: + # Dummy operands for the compile-time-elided reduction. + mailbox = sA + bar_reduce = bar_mma_epilog + + if cutlass.const_expr(self.publish is not None): + # Lamport publish: the reduced FP32 tile, M-major as the output tile. + sPub = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout((self.cta_m, self.cta_n)), + alignment=_PUBLISH_ALIGN_BYTES, + ) + else: + # Dummy operand for the compile-time-elided publish. + sPub = sA + + if warp_idx == 0: + with cute.arch.elect_one(): + for i in range(stages): + cute.arch.mbarrier_init(bar_full + i, 2) + cute.arch.mbarrier_init(bar_empty + i, 1) + cute.arch.mbarrier_init(bar_tma_epilog, 32) + cute.arch.mbarrier_init(bar_mma_epilog, 1) + cute.arch.mbarrier_init(bar_tmem_alloc, 160) + + if cutlass.const_expr(self.split_k > 1): + # Owner arrival plus peer transaction-byte credits. + cute.arch.mbarrier_init(bar_reduce, 1) + + cute.arch.mbarrier_init_fence() + if cutlass.const_expr(self.split_k > 1): + # Publish peer barriers before cross-CTA stores. + cute.arch.cluster_arrive_relaxed() + else: + cute.arch.barrier() + + # Host validation guarantees an equal, tail-free K partition. + k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k + k_tile_start = split_rank * k_tile_count + + if cutlass.const_expr(self.split_k > 1): + cute.arch.cluster_wait() + + # Warp 3 is idle; warps 4-7 run the epilogue. + if warp_idx == 0: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), + sA, + cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), + k_tile_start, + k_tile_count, + True, + ) + elif warp_idx == 1: + if cutlass.const_expr(self.hnorm_b is not None): + self.hnorm_b_warp(bar_full, bar_empty, sB, mB, mBias, k_tile_start) + else: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), + sB, + cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), + k_tile_start, + k_tile_count, + False, + ) + elif warp_idx == 2: + self.mma_warp( + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma, + sA, + sB, + tmem_base_ptr, + acc_layout, + self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), + k_tile_count, + ) + elif warp_idx >= 4: + self.epilog_warp( + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout, + cute.local_tile(mC, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.local_tile(mBias, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.arch.thread_idx()[0] - 128, + mC.element_type, + utils.LayoutEnum.from_tensor(mC), + mailbox, + bar_reduce, + split_rank, + sPub, + mC, + mLamport, + mFlags, + bidx, + n_idx, + ) + + @cute.experimental.jit + def dma_warp( + self, + bar_full, + bar_empty, + bar_tma_epilog, + g_tile: cute.Tensor, + s_tile: cute.Tensor, + cta_v_map: cute.Layout, + k_tile_start: cutlass.Int32, + k_tile_count: cutlass.Int32, + is_a: cutlass.Constexpr, + ): + stages = self.num_ab_stage + if cutlass.const_expr(not is_a and self.use_pdl): + cute.arch.griddepcontrol_wait() + + empty_phase = cutlass.Int32(1) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + bar_full + stage, + cute.size_in_bytes( + s_tile.element_type, + cute.slice_(s_tile.layout, (None, None, None, 0)), + ), + ) + cute_ext.tma_load( + g_tile[None, None, k_tile_start + k_tile], + s_tile[None, None, None, stage], + (bar_full + stage).value, + cta_v_map=cta_v_map, + tma_operation_type=self.tma_op, + update_expect_tx=False, + ) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + if cutlass.const_expr( + is_a + and self.use_pdl + and not (self.publish is not None and self.publish.trigger_after_wait) + ): + cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(not is_a and self.has_bias): + cute.arch.mbarrier_arrive(bar_tma_epilog) + self._drain_producer(bar_empty, empty_phase, k_tile_count) + + @cute.experimental.jit + def hnorm_b_warp( + self, + bar_full, + bar_empty, + sB: cute.Tensor, + mX: cute.Tensor, + mW: cute.Tensor, + k_tile_start: cutlass.Int32, + ): + """B of HnormB: FlashInfer's RMSNorm of the activation row x with weight w, written into + this CTA's B stages where the TMA would have put the norm's output. + + FlashInfer's CuTe kernel (H = 6144) gives a row to 64 threads; thread t holds the 8-element + vectors t + 64 b (b = 0..11), sums their squares in order, and the row sum is the butterfly + sum of threads 0-31 plus that of threads 32-63. Lane l of this warp plays threads l and + l + 32: it holds the vectors l + 32 m (m = 0..23; even m are thread l's, odd m thread + l + 32's). The B tiles are 8 x 128 K-major with the K_SW128 layout: row 0 (the activation; + rows 1-7 are zero, as the TMA fills rows past M) is unswizzled, K columns [0, 64) at bytes + [0, 128) and [64, 128) at bytes [1024, 1152) of the stage.""" + hidden = self.hnorm_b.hidden + lane = cute.arch.lane_idx() + vec_per_tile = self.cta_k // _HNORM_B_VEC_ELEMENTS + k_tiles = hidden // self.cta_k // self.split_k + per_lane = k_tiles * vec_per_tile // _WARP_LANES + row_rounds = hidden // _HNORM_B_VEC_ELEMENTS // _WARP_LANES + stage_bytes = self.cta_n * self.cta_k * _AB_ELEMENT_BYTES + atom_bytes = stage_bytes // 2 + row_bytes = atom_bytes // self.cta_n + sb = sB.iterator.toint() + zero = Int32(0) + + # Rows 1-7 of the stages this CTA uses. + zero_vecs = (atom_bytes - row_bytes) // 16 + for s in cutlass.range_constexpr(k_tiles): + for r in cutlass.range_constexpr((2 * zero_vecs + _WARP_LANES - 1) // _WARP_LANES): + idx = lane + Int32(r * _WARP_LANES) + if idx < 2 * zero_vecs: + atom = idx // zero_vecs + _st_shared_v4( + sb + + Int32(s * stage_bytes + row_bytes) + + atom * atom_bytes + + (idx - atom * zero_vecs) * 16, + zero, + zero, + zero, + zero, + ) + + # This CTA's K slice is the vectors [v0, v0 + k_tiles * vec_per_tile); lane l produces the + # slice vectors v0 + l + 32 j. Their norm weights are constants: loaded before the wait. + v0 = k_tile_start * vec_per_tile + w_words = [] + for j in cutlass.range_constexpr(per_lane): + v = v0 + lane + Int32(j * _WARP_LANES) + w_words.append(_ld_global_cg_v4((mW.iterator + v * _HNORM_B_VEC_ELEMENTS).toint())) + + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + + x_words = [] + for m in cutlass.range_constexpr(row_rounds): + v = lane + Int32(m * _WARP_LANES) + x_words.append(_ld_global_cg_v4((mX.iterator + v * _HNORM_B_VEC_ELEMENTS).toint())) + xs_words = [] + for j in cutlass.range_constexpr(per_lane): + v = v0 + lane + Int32(j * _WARP_LANES) + xs_words.append(_ld_global_cg_v4((mX.iterator + v * _HNORM_B_VEC_ELEMENTS).toint())) + + acc_lo = Float32(0.0) + acc_hi = Float32(0.0) + for b in cutlass.range_constexpr(row_rounds // 2): + for q in cutlass.range_constexpr(4): + acc_lo = _add_bf16x2_squares(acc_lo, x_words[2 * b][q]) + for b in cutlass.range_constexpr(row_rounds // 2): + for q in cutlass.range_constexpr(4): + acc_hi = _add_bf16x2_squares(acc_hi, x_words[2 * b + 1][q]) + for offset in (1, 2, 4, 8, 16): + acc_lo = _add_rn(acc_lo, cute.arch.shuffle_sync_bfly(acc_lo, offset)) + acc_hi = _add_rn(acc_hi, cute.arch.shuffle_sync_bfly(acc_hi, offset)) + rstd = _rmsnorm_rstd(_add_rn(acc_lo, acc_hi), hidden, self.hnorm_b.eps) + + for j in cutlass.range_constexpr(per_lane): + local = lane + Int32(j * _WARP_LANES) + tile = local // vec_per_tile + col_vec = local - tile * vec_per_tile + col_atom = col_vec // (vec_per_tile // 2) + _st_shared_v4( + sb + + tile * stage_bytes + + col_atom * atom_bytes + + (col_vec - col_atom * (vec_per_tile // 2)) * 16, + _rmsnorm_bf16x2(xs_words[j][0], w_words[j][0], rstd), + _rmsnorm_bf16x2(xs_words[j][1], w_words[j][1], rstd), + _rmsnorm_bf16x2(xs_words[j][2], w_words[j][2], rstd), + _rmsnorm_bf16x2(xs_words[j][3], w_words[j][3], rstd), + ) + # The tensor core reads the stages through the async proxy. + _fence_proxy_async_shared() + cute.arch.sync_warp() + with cute.arch.elect_one(): + for s in cutlass.range_constexpr(k_tiles): + cute.arch.mbarrier_arrive(bar_full + s) + # As the TMA producer: wait until the MMA has released every stage this CTA filled. + self._drain_producer(bar_empty, Int32(1 ^ (k_tiles // self.num_ab_stage)), Int32(k_tiles)) + + @cute.experimental.jit + def _drain_producer( + self, + bar_empty, + empty_phase: cutlass.Int32, + k_tile_count: cutlass.Int32, + ): + stages = self.num_ab_stage + for tail in cutlass.range(stages, unroll=1): + stage = (tail + k_tile_count) % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + @cute.experimental.jit + def mma_warp( + self, + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma: cute.TiledMma, + sA: cute.Tensor, + sB: cute.Tensor, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + mma_inst_tile_k: cutlass.Constexpr, + k_tile_count: cutlass.Int32, + ): + num_tmem_cols = 256 + cute.arch.alloc_tmem(num_tmem_cols, tmem_base_ptr, is_two_cta=False) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False) + + tmem_ptr = cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr) + accumulator = cute.make_tensor(tmem_ptr, acc_layout)[None, None, None, 0] + mma_atom = cute.make_mma_atom(tiled_mma.op) + full_phase = cutlass.Int32(0) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % self.num_ab_stage + cute.arch.mbarrier_wait(bar_full + stage, full_phase) + for k_block in range(mma_inst_tile_k): + if k_block == 0: + mma_atom.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + else: + mma_atom.set(tcgen05.Field.ACCUMULATE, True) + cute_ext.dot( + mma_atom, + cute.append_ones(sA[None, None, k_block, stage], up_to_rank=3), + cute.append_ones(sB[None, None, k_block, stage], up_to_rank=3), + accumulator, + ) + with cute.arch.elect_one(): + tcgen05.commit(bar_empty + stage, None, self.cta_group) + if stage == self.num_ab_stage - 1: + full_phase = full_phase ^ 1 + + with cute.arch.elect_one(): + tcgen05.commit(bar_mma_epilog, None, self.cta_group) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 1) + cute.arch.dealloc_tmem(tmem_ptr, num_tmem_cols, is_two_cta=False) + + @cute.experimental.jit + def epilog_warp( + self, + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + gD_tile: cute.Tensor, + gBias_tile: cute.Tensor, + epi_tid: cutlass.Int32, + c_dtype: cutlass.Constexpr, + d_layout: cutlass.Constexpr, + mailbox, + bar_reduce, + split_rank: cutlass.Int32, + sPub, + mC: cute.Tensor, + mLamport: cute.Tensor, + mFlags: cute.Tensor, + bidx: cutlass.Int32, + n_idx: cutlass.Int32, + ): + lamport_buffer = Int64(0) + if cutlass.const_expr(self.publish is not None): + # The flags are read after the grid dependency wait: the previous all-reduce, which + # rotated them, precedes this grid's primary. Every CTA's epilogue warps wait; with + # trigger_after_wait they also release the dependents here. + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(self.publish.trigger_after_wait): + cute.arch.griddepcontrol_launch_dependents() + lamport_buffer = _lamport_current_buffer( + mFlags.iterator.toint(), mLamport.iterator.toint() + ) + + # Wait until MMA publishes the TMEM base pointer. + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) + + acc_view = cute.make_tensor( + cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr), + acc_layout, + )[((None, None), 0, 0, 0)] + + epi_tile = (self.cta_m, self.cta_n) + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( + sm100_utils.get_tmem_load_op( + (self.cta_m, self.cta_n, self.cta_k), + d_layout, + c_dtype, + self.acc_dtype, + epi_tile, + False, + ), + acc_view, + ) + gD_epi = cute.flat_divide(gD_tile, epi_tile) + + # Match each epilogue thread's TMEM partition in RMEM. + rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid) + rAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rD = cute_ext.allocate( + c_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + thr_t2r = tiled_copy_t2r.get_slice(epi_tid) + + if cutlass.const_expr(self.has_bias): + bias_dtype = gBias_tile.element_type + rBias = cute_ext.allocate( + bias_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rBiasAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + if split_rank == OWNER_RANK: + cute.arch.mbarrier_wait(bar_tma_epilog, 0) + cute_ext.partition_and_copy( + cute.make_tiled_copy_D( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), bias_dtype), + tiled_copy_t2r, + ).get_slice(epi_tid), + cute.flat_divide(gBias_tile, epi_tile)[None, None, 0, 0], + rBias, + ) + rBiasAcc.store(rBias.load().to(self.acc_dtype)) + + cute.arch.mbarrier_wait(bar_mma_epilog, 0) + cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) + # Make tcgen05.ld visible before TMEM release and RMEM use. + cute.arch.fence_view_async_tmem_load() + cute.arch.mbarrier_arrive(bar_tmem_alloc) + + # Peers publish FP32 partials; only rank 0 reduces and stores. + if cutlass.const_expr(self.split_k > 1): + assert cute.size(rmem_layout) == self.mailbox_elements // ( + (self.split_k - 1) * self.epilog_threads + ) + values_per_thread = cutlass.const_expr(cute.size(rmem_layout)) + values_per_peer = cutlass.const_expr(self.epilog_threads * values_per_thread) + if split_rank != OWNER_RANK: + for value_idx in cutlass.range_constexpr(0, values_per_thread, 4): + _store_shared_remote_v4( + rAcc[value_idx], + rAcc[value_idx + 1], + rAcc[value_idx + 2], + rAcc[value_idx + 3], + mailbox.iterator + + (split_rank - Int32(1)) * values_per_peer + + epi_tid * values_per_thread + + value_idx, + bar_reduce, + Int32(OWNER_RANK), + ) + else: + if epi_tid == 0: + cute.arch.mbarrier_arrive_and_expect_tx( + bar_reduce, self.expected_transaction_bytes + ) + cute.arch.mbarrier_wait(bar_reduce, 0) + for peer in cutlass.range_constexpr(self.split_k - 1): + for value_idx in cutlass.range_constexpr(values_per_thread): + rAcc[value_idx] = ( + rAcc[value_idx] + + mailbox[ + peer * values_per_peer + epi_tid * values_per_thread + value_idx + ] + ) + + if split_rank == OWNER_RANK: + if cutlass.const_expr(self.publish is not None): + # Stage the tile by output coordinate, then store eight columns per thread. + cute_ext.partition_and_copy(thr_t2r, rAcc, sPub) + cute.arch.barrier( + barrier_id=_PUBLISH_BARRIER_ID, number_of_threads=self.epilog_threads + ) + self._publish_tile(sPub, mC, lamport_buffer, bidx, n_idx, epi_tid) + else: + if cutlass.const_expr(self.has_bias): + rAcc.store(rAcc.load() + rBiasAcc.load()) + + rD.store(rAcc.load().to(c_dtype)) + # Preserve TMEM coordinates; the copy predicates output tails. + cute_ext.partition_and_copy(thr_t2r, rD, gD_epi[None, None, 0, 0]) + + # The reduction mbarrier covers remote stores; no cluster barrier needed. + + @cute.experimental.jit + def _publish_tile( + self, + sPub: cute.Tensor, + mC: cute.Tensor, + lamport_buffer: Int64, + bidx: cutlass.Int32, + n_idx: cutlass.Int32, + epi_tid: cutlass.Int32, + ): + """Store the staged tile into the Lamport buffer: eight consecutive output columns of one + row per thread, rounded to BF16 as C would be, -0.0 as +0.0, one 16-byte store each at + [token][rank][column] (the one-shot all-reduce's slot of this rank).""" + vecs_m = cutlass.const_expr(self.cta_m // _PUBLISH_VEC_ELEMENTS) + num_vecs = cutlass.const_expr(vecs_m * self.cta_n) + n_feat = cute.size(mC, mode=[0]) + n_tok = cute.size(mC, mode=[1]) + for it in cutlass.range_constexpr( + (num_vecs + self.epilog_threads - 1) // self.epilog_threads + ): + vec = epi_tid + it * self.epilog_threads + row = vec // vecs_m + m8 = (vec - row * vecs_m) * _PUBLISH_VEC_ELEMENTS + token = n_idx * self.cta_n + row + feature = bidx * self.cta_m + m8 + if vec < num_vecs: + if token < n_tok: + if feature < n_feat: + words = [] + for j in cutlass.range_constexpr(_PUBLISH_VEC_ELEMENTS // 2): + words.append( + _cvt_bf16x2_lamport( + sPub[m8 + 2 * j, row], sPub[m8 + 2 * j + 1, row] + ) + ) + slot = Int64(token * self.publish.world_size + self.publish.rank) + offset = (slot * Int64(n_feat) + Int64(feature)) * 2 + _store_global_v4( + lamport_buffer + offset, words[0], words[1], words[2], words[3] + ) + + +_SUPPORTED_TORCH_DTYPES = (_torch.bfloat16, _torch.float16) + + +@cute.experimental.jit +def _bmm_no_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + stream: _cuda.CUstream, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + # C also stands in for the dead bias and Lamport operands. + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + c, + c, + stream, + ) + + +@cute.experimental.jit +def _bmm_hnorm_b( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + norm_weight: cute.Tensor, + stream: _cuda.CUstream, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + # HnormB: b is the raw activation and the norm weight takes the dead bias operand's place (with + # C's layout: only its address is read); C stands in for the dead Lamport operands. + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(norm_weight.iterator, c.layout), + c, + c, + stream, + ) + + +@cute.experimental.jit +def _bmm_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + stream: _cuda.CUstream, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + # C also stands in for the dead Lamport operands. + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(bias.iterator, cute.select(bias.layout, mode=[1, 2, 0])), + c, + c, + stream, + ) + + +@cute.experimental.jit +def _bmm_publish( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + lamport: cute.Tensor, + flags: cute.Tensor, + stream: _cuda.CUstream, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + # C (not written) also stands in for the dead bias. + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + c, + lamport, + flags, + stream, + ) + + +def _from_dlpack_dynamic(tensor, leading_dim: int, assumed_align: int = 32): + # Detached: DLPack exports no tensor that requires grad, which module parameters (the LM head + # weight, for one) may still do; the view shares the storage. + return from_dlpack(tensor.detach(), assumed_align=assumed_align).mark_layout_dynamic( + leading_dim=leading_dim + ) + + +def _detect_leading_dim(tensor: _torch.Tensor) -> int: + # Ignore synthetic batch stride, including 1x1. + for dim, stride in enumerate(tensor.stride()[1:], start=1): + if stride == 1: + return dim + raise ValueError("tensor has no stride-1 dimension") + + +def _make_layout_tensor( + shape: tuple[int, ...], dtype: _torch.dtype, leading_dim: int +) -> _torch.Tensor: + permutation = [dim for dim in range(len(shape)) if dim != leading_dim] + [leading_dim] + return _torch.empty( + tuple(shape[dim] for dim in permutation), dtype=dtype, device="cuda" + ).permute([permutation.index(dim) for dim in range(len(shape))]) + + +def _make_compile_repr_tensors( + dtype: _torch.dtype, + has_bias: bool, + a_leading: int, + b_leading: int, + c_leading: int, +): + m, n, k, batch = 64, 8, _CTA_K, 1 + tensors = tuple( + _from_dlpack_dynamic(_make_layout_tensor(shape, dtype, leading_dim), leading_dim) + for shape, leading_dim in ( + ((batch, n, k), a_leading), + ((batch, k, m), b_leading), + ((batch, n, m), c_leading), + ) + ) + if not has_bias: + return (*tensors, None) + return ( + *tensors, + _from_dlpack_dynamic( + _torch.empty((n,), dtype=dtype, device="cuda").as_strided( + size=(batch, n, m), stride=(0, 1, 0) + ), + 1, + 2, + ), + ) + + +def _to_cute_swap(a, b, out, bias): + a_swap = b.unsqueeze(0).transpose(-2, -1) + b_swap = a.unsqueeze(0).transpose(-2, -1) + c_swap = out.unsqueeze(0).transpose(-2, -1) + leading_dims = tuple(_detect_leading_dim(tensor) for tensor in (a_swap, b_swap, c_swap)) + cute_tensors = tuple( + _from_dlpack_dynamic(tensor, leading_dim) + for tensor, leading_dim in zip((a_swap, b_swap, c_swap), leading_dims, strict=True) + ) + if bias is None: + return (*cute_tensors, None, leading_dims) + return ( + *cute_tensors, + _from_dlpack_dynamic( + bias.as_strided(size=(1, c_swap.shape[1], c_swap.shape[2]), stride=(0, 1, 0)), + 1, + 2, + ), + leading_dims, + ) + + +@functools.cache +def _get_compiled_splitk_kernel( + dtype, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + leading_dims: tuple[int, int, int], + reserve_full_smem: bool = True, + publish: LamportPublish | None = None, +): + if dtype not in _SUPPORTED_TORCH_DTYPES: + raise ValueError(f"split-K dense GEMM supports {_SUPPORTED_TORCH_DTYPES}; got {dtype}") + + kernel = SplitKDenseGemmKernel( + tactic=tactic, + use_pdl=use_pdl, + has_bias=has_bias, + reserve_full_smem=reserve_full_smem, + publish=publish, + ) + compile_tensors = _make_compile_repr_tensors(dtype, has_bias, *leading_dims) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + if publish is not None: + lamport = _from_dlpack_dynamic(_torch.empty((64,), dtype=dtype, device="cuda"), 0) + flags = _from_dlpack_dynamic(_torch.empty((16,), dtype=_torch.int32, device="cuda"), 0) + compiled = cute_ext.compile( + _bmm_publish, kernel, *compile_tensors[:3], lamport, flags, stream + ) + elif has_bias: + compiled = cute_ext.compile(_bmm_bias, kernel, *compile_tensors, stream) + else: + compiled = cute_ext.compile(_bmm_no_bias, kernel, *compile_tensors[:3], stream) + return compiled + + +def operands_overlap(a: _torch.Tensor, b: _torch.Tensor) -> bool: + """Whether the byte ranges that ``a`` and ``b`` span on one device intersect.""" + if a.numel() == 0 or b.numel() == 0 or a.device != b.device: + return False + + def span(t: _torch.Tensor) -> int: + extent = sum((size - 1) * abs(stride) for size, stride in zip(t.shape, t.stride())) + return (extent + 1) * t.element_size() + + a_lo, b_lo = a.data_ptr(), b.data_ptr() + return a_lo < b_lo + span(b) and b_lo < a_lo + span(a) + + +def _validate_runtime_tensors(a, b, bias, out) -> tuple[int, int, int]: + tensors = (a, b, out) + ((bias,) if bias is not None else ()) + if any(not isinstance(tensor, _torch.Tensor) for tensor in tensors): + raise ValueError("a, b, out, and bias must be torch tensors") + if a.ndim != 2 or b.ndim != 2 or out.ndim != 2: + raise ValueError("split-K dense GEMM accepts only 2D tensors") + if a.device.type != "cuda" or any(tensor.device != a.device for tensor in tensors): + raise ValueError("all tensors must be on the same CUDA device") + if a.dtype not in _SUPPORTED_TORCH_DTYPES or any(tensor.dtype != a.dtype for tensor in tensors): + raise ValueError("a, b, out, and bias must share BF16 or FP16 dtype") + if any(not (tensor.is_contiguous() or tensor.t().is_contiguous()) for tensor in (a, b, out)): + raise ValueError("a, b, and out must be dense row-major or column-major matrices") + if any(tensor.data_ptr() % 32 for tensor in (a, b, out)): + raise ValueError("a, b, and out must be 32-byte aligned") + # cute_ext.tma_load is used for a and b (not out, which uses partition_and_copy). + # TMA requires the non-innermost byte stride to be a multiple of 16 bytes. + # For a dense 2D tensor this is max(stride(0), stride(1)) * element_size. + _TMA_STRIDE_ALIGN = 16 + for tensor, tname in ((a, "a"), (b, "b")): + outer_stride_bytes = max(tensor.stride(0), tensor.stride(1)) * tensor.element_size() + if outer_stride_bytes % _TMA_STRIDE_ALIGN: + raise ValueError( + f"tensor '{tname}' outer byte stride {outer_stride_bytes} is not a " + f"multiple of {_TMA_STRIDE_ALIGN}; TMA load requires " + f"{_TMA_STRIDE_ALIGN}-byte-aligned strides " + f"(shape={tuple(tensor.shape)}, strides={tensor.stride()})." + ) + + m, k = a.shape + if b.shape[0] != k: + raise ValueError(f"incompatible shapes: a is {tuple(a.shape)}, b is {tuple(b.shape)}") + n = b.shape[1] + if out.shape != (m, n): + raise ValueError(f"out must have shape {(m, n)}, got {tuple(out.shape)}") + if bias is not None and (bias.ndim != 1 or bias.shape[0] != n or not bias.is_contiguous()): + raise ValueError( + f"bias must be contiguous with shape {(n,)}, " + f"got shape {tuple(bias.shape)} and stride {bias.stride()}" + ) + + return m, n, k + + +def run_splitk_dense( + a: _torch.Tensor, + b: _torch.Tensor, + bias: _torch.Tensor | None, + out: _torch.Tensor, + pdl: bool, + tactic: SplitKTactic, + reserve_full_smem: bool = True, +) -> _torch.Tensor: + """Run ``A[M,K] @ B[K,N]`` with the split-K BF16/FP16 GEMM kernel. + + Args: + a: Input tensor of shape ``[M, K]``, row-major BF16/FP16. + b: Weight tensor of shape ``[K, N]``, column-major BF16/FP16 + (i.e. the transpose of a ``[N, K]`` contiguous weight matrix). + bias: Optional bias of shape ``[N]``, or ``None``. + out: Pre-allocated output tensor of shape ``[M, N]``, BF16/FP16. + pdl: Whether to use Programmatic Dependent Launch. + tactic: The ``SplitKTactic`` specifying tile and pipeline parameters. + reserve_full_smem: Reserve the full shared-memory capacity per CTA (one CTA per + SM). ``False`` reserves only ``tactic_launch_smem_bytes(tactic)``; same kernel + and result. + + Returns: + The ``out`` tensor, filled with the GEMM result. + """ + validate_tactic(tactic, *_validate_runtime_tensors(a, b, bias, out)) + has_bias = bias is not None + cute_tensors = _to_cute_swap(a, b, out, bias) + compiled = _get_compiled_splitk_kernel( + dtype=a.dtype, + tactic=tactic, + use_pdl=pdl, + has_bias=has_bias, + leading_dims=cute_tensors[4], + reserve_full_smem=reserve_full_smem, + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + if has_bias: + compiled(*cute_tensors[:4], stream) + else: + compiled(*cute_tensors[:3], stream) + return out + + +def publish_tactic( + tactic: SplitKTactic, *, smem_capacity: int = _SMEM_CAPACITY_BYTES +) -> SplitKTactic: + """``tactic`` with at most as many A/B stages as fit beside the publish staging tile (the + same results; only the pipeline depth changes).""" + stages = tactic.ab_stages + while stages > 1 and _publish_smem_bytes(tactic, stages) > smem_capacity: + stages -= 1 + return dataclasses.replace(tactic, ab_stages=stages) + + +def publish_supported( + m: int, + n: int, + tactic: SplitKTactic, + *, + dtype: _torch.dtype = _torch.bfloat16, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> bool: + """Whether ``run_splitk_publish`` serves an ``[m, n]`` BF16 output with ``tactic``: 16-byte + column groups and the staging tile beside the A/B stages.""" + return ( + dtype == _torch.bfloat16 + and 1 <= m <= MAX_M + and n > 0 + and n % _PUBLISH_VEC_ELEMENTS == 0 + and tactic.mma_m % _PUBLISH_VEC_ELEMENTS == 0 + and _publish_smem_bytes(tactic, tactic.ab_stages) <= smem_capacity + ) + + +def run_splitk_publish( + a: _torch.Tensor, + b: _torch.Tensor, + out: _torch.Tensor, + lamport: _torch.Tensor, + flags: _torch.Tensor, + pdl: bool, + tactic: SplitKTactic, + publish: LamportPublish, + reserve_full_smem: bool = True, +) -> None: + """Run ``A[M,K] @ B[K,N]`` (no bias) and publish the BF16 result into an MNNVL one-shot + all-reduce workspace instead of writing ``out``. + + Args: + a, b, tactic, pdl, reserve_full_smem: as for ``run_splitk_dense``. + out: The ``[M, N]`` output ``run_splitk_dense`` would write; it only gives the output + layout and is not written. + lamport: The workspace's multicast view (BF16, 1-D, its NUM_LAMPORT_BUFFERS buffers). + flags: The workspace's buffer flags viewed as int32. + publish: World size and rank of the all-reduce group, and the trigger placement. + + Every value lands at ``[token][publish.rank][column]`` of the buffer ``flags`` marks current, on + every rank of the group, as the shard store of the one-shot kernel would place it; the + all-reduce of this output must then skip its own shard load and store. + """ + m, n, k = _validate_runtime_tensors(a, b, None, out) + validate_tactic(tactic, m, n, k) + if not publish_supported(m, n, tactic, dtype=a.dtype): + raise ValueError(f"Lamport publish does not serve M={m}, N={n} with {tactic}") + if ( + lamport.dtype != a.dtype + or lamport.dim() != 1 + or flags.dtype != _torch.int32 + or flags.dim() != 1 + ): + raise ValueError("lamport must be a 1-D view of the output dtype, flags a 1-D int32 view") + cute_tensors = _to_cute_swap(a, b, out, None) + compiled = _get_compiled_splitk_kernel( + dtype=a.dtype, + tactic=tactic, + use_pdl=pdl, + has_bias=False, + leading_dims=cute_tensors[4], + reserve_full_smem=reserve_full_smem, + publish=publish, + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + compiled( + *cute_tensors[:3], + _from_dlpack_dynamic(lamport, 0), + _from_dlpack_dynamic(flags, 0), + stream, + ) + + +@functools.cache +def _get_compiled_hidden_half_kernel( + tactic: SplitKTactic, + use_pdl: bool, + leading_dims: tuple[int, int, int], + norm_eps: float | None, +): + """The split-K GEMM with an FP32 C; with ``norm_eps``, the ``HnormB`` variant that takes the + norm weight as a fourth operand.""" + hnorm_b = None if norm_eps is None else HnormB(HNORM_B_HIDDEN, norm_eps) + kernel = SplitKDenseGemmKernel(tactic=tactic, use_pdl=use_pdl, has_bias=False, hnorm_b=hnorm_b) + # The representative tensors of _make_compile_repr_tensors, with an FP32 C. + m, n, k, batch = 64, 8, _CTA_K, 1 + compile_tensors = tuple( + _from_dlpack_dynamic(_make_layout_tensor(shape, dtype, leading_dim), leading_dim) + for shape, dtype, leading_dim in ( + ((batch, n, k), _torch.bfloat16, leading_dims[0]), + ((batch, k, m), _torch.bfloat16, leading_dims[1]), + ((batch, n, m), _torch.float32, leading_dims[2]), + ) + ) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + if hnorm_b is None: + return cute_ext.compile(_bmm_no_bias, kernel, *compile_tensors, stream) + norm_weight = _from_dlpack_dynamic(_torch.empty((k,), dtype=_torch.bfloat16, device="cuda"), 0) + return cute_ext.compile(_bmm_hnorm_b, kernel, *compile_tensors, norm_weight, stream) + + +def hidden_half_norm_supported( + x: _torch.Tensor, weight: _torch.Tensor, norm_weight: _torch.Tensor, tactic: SplitKTactic +) -> bool: + """Whether ``run_splitk_hidden_half`` can normalize ``x`` itself (``HnormB``): one 16-byte + aligned BF16 row of ``HNORM_B_HIDDEN`` with unit stride, a contiguous 16-byte aligned BF16 + norm weight of that size, a weight of that K, and a tactic of ``hnorm_b_supported``.""" + return ( + x.dim() == 2 + and x.shape[0] == 1 + and x.dtype == _torch.bfloat16 + and x.shape[1] == HNORM_B_HIDDEN + and x.stride(1) == 1 + and x.data_ptr() % 16 == 0 + and norm_weight.dtype == _torch.bfloat16 + and tuple(norm_weight.shape) == (x.shape[1],) + and norm_weight.is_contiguous() + and norm_weight.data_ptr() % 16 == 0 + and weight.dim() == 2 + and weight.shape[1] == x.shape[1] + and hnorm_b_supported(tactic, x.shape[1]) + ) + + +def run_splitk_hidden_half( + x: _torch.Tensor, + weight: _torch.Tensor, + out: _torch.Tensor, + pdl: bool, + tactic: SplitKTactic, + norm_weight: _torch.Tensor | None = None, + norm_eps: float | None = None, +) -> _torch.Tensor: + """``x @ weight.T`` into an FP32 ``out``, or with ``norm_weight`` ``rmsnorm(x) @ weight.T`` + with the norm computed in the kernel (``HnormB``). + + Args: + x: Activation [rows, K] BF16 with unit inner stride; with ``norm_weight`` the one raw + row whose FlashInfer RMSNorm (``norm_weight``, ``norm_eps``) the kernel computes. + weight: [N, K] BF16, contiguous. + out: [rows, N] FP32, contiguous: the split-K FP32 sum, never rounded to BF16. + pdl: Whether to use Programmatic Dependent Launch. + tactic: The ``SplitKTactic``; with ``norm_weight`` one of ``hnorm_b_supported``. + norm_weight: [K] BF16 RMSNorm weight, or None for an ``x`` that is already normalized. + norm_eps: The RMSNorm epsilon (with ``norm_weight``). + + Returns: + ``out``. + + With ``norm_weight`` the result is bit-identical to running FlashInfer's RMSNorm first and + this GEMM on its BF16 output. + """ + if x.dim() != 2 or weight.dim() != 2 or out.dim() != 2: + raise ValueError("x, weight and out must be 2D") + rows, k = x.shape + n = weight.shape[0] + if ( + x.dtype != _torch.bfloat16 + or weight.dtype != _torch.bfloat16 + or out.dtype != _torch.float32 + or weight.shape[1] != k + or tuple(out.shape) != (rows, n) + or not weight.is_contiguous() + or not out.is_contiguous() + or x.stride(1) != 1 + or any(t.data_ptr() % 32 for t in (x, weight, out)) + ): + raise ValueError( + "run_splitk_hidden_half needs BF16 x [rows, K] with unit inner stride, a contiguous " + "BF16 weight [N, K] and a contiguous FP32 out [rows, N], all 32-byte aligned; got " + f"x {tuple(x.shape)} {x.dtype}, weight {tuple(weight.shape)} {weight.dtype}, " + f"out {tuple(out.shape)} {out.dtype}" + ) + validate_tactic(tactic, rows, n, k) + if norm_weight is not None and ( + norm_eps is None or not hidden_half_norm_supported(x, weight, norm_weight, tactic) + ): + raise ValueError( + f"the in-kernel norm does not serve x {tuple(x.shape)}, weight " + f"{tuple(weight.shape)} with {tactic}" + ) + cute_tensors = _to_cute_swap(x, weight.t(), out, None) + compiled = _get_compiled_hidden_half_kernel( + tactic, pdl, cute_tensors[4], None if norm_weight is None else float(norm_eps) + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(x.device).cuda_stream) + if norm_weight is None: + compiled(*cute_tensors[:3], stream) + else: + compiled(*cute_tensors[:3], _from_dlpack_dynamic(norm_weight.detach(), 0), stream) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk_fused.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk_fused.py new file mode 100644 index 000000000000..89b40efdca3b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk_fused.py @@ -0,0 +1,2021 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-M split-K BF16 GEMM with fused epilogues for the MLA q_b projection (SM10x, CuTe DSL). + +``SplitKFusedEpilogueGemmKernel`` runs the mainloop and the cluster split-K reduction of +``_low_m_splitk.SplitKDenseGemmKernel`` unchanged. Only the owner CTA's epilogue differs: +it stages the reduced FP32 tile in shared memory and then stores four consecutive output +columns per thread, so per-row and per-column-pair operations do not depend on the TMEM load +layout. Options (``SplitKEpilogue``): + +``row_rms`` + B (the public activation ``a``) is the raw input x of an RMSNorm whose weight is folded + into the GEMM weight (``W * gamma``). The owner's epilogue warps compute every row's mean + square from global memory while the mainloop runs, and the row's outputs are scaled by + ``rsqrt(mean + eps)``: ``(x * rstd) @ (W * gamma)^T``, which is ``RMSNorm(x) @ W^T`` up to + rounding. +``rms_out`` + Also write ``RMSNorm(x) = x * rstd * gamma`` (BF16 [M, K]) for other readers of the + normalized activation. The owner CTAs of a row tile split its K columns. +``fp8_rope`` + Write FP8 E4M3 (times a device scale) instead of BF16, after GPT-J RoPE on the last + ``rope_dim`` columns of every ``head_dim``-wide head: the MLA generation Q layout + [tokens, heads, kv_lora_rank + qk_rope_head_dim] that ``mla_rope_generation`` rotates + and quantizes, with its position rule (``kv_lens[seq] - seq_len + token_in_seq``), its + cos/sin table layout and its saturating E4M3 conversion. The BF16 rounding the unfused + path applies before the rotation and before the quantization is skipped. + +``SplitKIndexerQTailGemmKernel`` (``run_splitk_idx_qtail``) is the ``row_rms`` kernel of the DSA +indexer wq_b whose owner epilogue finishes the indexer Q (``IndexerQTail``) instead of storing it: +it rounds every output to BF16 as the BF16 store does, then applies the per-(token, head) tail of +the fused DSA indexer prologue to those values with the same operations in the same order (RoPE, +BF16 rounding, the per-head E4M3 scale, the E4M3 Q, the scale, and the attention weights times +the scale). Head-wide tiles finish a head per warp; half-head tiles run +in clusters of the two halves, whose owners exchange their row amax through DSMEM, so the GEMM +keeps its 64-row tactic and its K split. +""" + +from __future__ import annotations + +import dataclasses +import functools +import math + +import cuda.bindings.driver as _cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import torch as _torch +from cutlass import Float32, Int32, Int64 +from cutlass._mlir.dialects import llvm +from cutlass.cute import experimental as cute_ext +from cutlass.cutlass_dsl import T, dsl_user_op + +from ._low_m_splitk import ( + _AB_BUFFER_ALIGN_BYTES, + _CTA_K, + _LAUNCH_SMEM_HEADROOM_BYTES, + _MAILBOX_ALIGN_BYTES, + _MBARRIER_BYTES, + _SMEM_CAPACITY_BYTES, + _TMEM_POINTER_BYTES, + OWNER_RANK, + SplitKDenseGemmKernel, + SplitKTactic, + _align_up, + _detect_leading_dim, + _from_dlpack_dynamic, + _smem_bytes, + _store_shared_remote_v4, + default_tactic, + device_limits, + splitk_default_tactic, + validate_tactic, +) + +__all__ = [ + "IndexerQTail", + "Q_B_PLANS", + "SplitKEpilogue", + "fused_epilogue_tactic", + "idx_qtail_plan", + "idx_qtail_tactic", + "q_b_plan", + "run_splitk_fused", + "run_splitk_idx_qtail", +] + +#: Alignment of the epilogue's shared-memory staging buffers. +_EPI_ALIGN_BYTES = 16 + +#: Named barrier of the four epilogue warps (barrier 0 is the CTA-wide one). +_EPI_BARRIER_ID = 1 + +#: BF16 elements per 16-byte vector. +_BF16_PER_VEC = 8 + +#: Lanes of a warp. +_WARP_LANES = 32 + + +@dataclasses.dataclass(frozen=True) +class SplitKEpilogue: + """Fused epilogue options; every field is a compile-time constant of the kernel.""" + + row_rms: bool = False + rms_k: int = 0 + rms_eps: float = 0.0 + rms_out: bool = False + fp8_rope: bool = False + head_dim: int = 0 + rope_dim: int = 0 + seq_len: int = 1 + + def validate(self) -> None: + if not (self.row_rms or self.fp8_rope): + raise ValueError("a fused epilogue needs row_rms or fp8_rope") + if self.row_rms and (self.rms_k <= 0 or self.rms_k % _BF16_PER_VEC or self.rms_eps <= 0): + raise ValueError(f"row_rms needs a positive rms_k divisible by 8 and eps > 0: {self}") + if self.rms_out and not self.row_rms: + raise ValueError("rms_out needs row_rms") + if self.fp8_rope and not ( + 0 < self.rope_dim < self.head_dim + and self.rope_dim % 4 == 0 + and (self.head_dim - self.rope_dim) % 4 == 0 + and self.seq_len >= 1 + ): + raise ValueError(f"invalid fp8_rope geometry: {self}") + + +def _fused_smem_bytes(tactic: SplitKTactic, ab_stages: int) -> int: + """Shared memory of the fused kernel: the base layout, then the FP32 tile and row scales.""" + cursor = _align_up(_smem_bytes(tactic, ab_stages), _EPI_ALIGN_BYTES) + cursor += tactic.mma_m * tactic.mma_n * 4 + return _align_up(cursor, _EPI_ALIGN_BYTES) + tactic.mma_n * 4 + + +def fused_epilogue_tactic(tactic: SplitKTactic) -> SplitKTactic: + """``tactic`` with at most as many A/B stages as fit beside the epilogue staging buffers.""" + stages = tactic.ab_stages + while stages > 1 and _fused_smem_bytes(tactic, stages) > _SMEM_CAPACITY_BYTES: + stages -= 1 + return dataclasses.replace(tactic, ab_stages=stages) + + +@dataclasses.dataclass(frozen=True) +class IndexerQTail: + """The DSA indexer Q tail of ``run_splitk_idx_qtail``; every field is a compile-time constant. + + Per (token, head) of the BF16-rounded wq_b output, as ``_dsa_indexer_prologue_fused_kernel``: + RoPE on the first ``rope_dim`` columns (GPT-J pairs with ``interleave``, NeoX halves otherwise) + with the cos/sin row of the token's position, rounded to BF16; the head's E4M3 scale + max(amax, 1e-12) / 448, rounded up to a power of two with ``ue8m0``; the E4M3 Q (times the + MUFU reciprocal of the scale); the scale; and the attention weight times the scale times + ``weight_scale``. + """ + + n_heads: int + head_dim: int = 128 + rope_dim: int = 64 + interleave: bool = True + ue8m0: bool = True + weight_scale: float = 1.0 + + def validate(self) -> None: + if not ( + self.n_heads > 0 + and self.head_dim % 8 == 0 + and 0 < self.rope_dim <= self.head_dim // 2 + and self.rope_dim % 8 == 0 + and self.weight_scale > 0 + ): + raise ValueError(f"invalid indexer Q tail geometry: {self}") + + +def _qtail_pair(tactic: SplitKTactic, qtail: IndexerQTail) -> bool: + """Whether a tile holds half a head (clusters of the two halves) rather than a whole head.""" + if tactic.mma_m == qtail.head_dim: + return False + if 2 * tactic.mma_m == qtail.head_dim: + return True + raise ValueError(f"tile of {tactic.mma_m} columns is neither a head nor half a head: {qtail}") + + +def _qtail_smem_bytes(tactic: SplitKTactic, ab_stages: int, qtail: IndexerQTail) -> int: + """Shared memory of the indexer Q tail kernel: the fused layout, then the cos/sin rows and the + attention weights of the tile's tokens, then (half-head tiles) both halves' row amax and their + mbarrier.""" + cursor = _align_up(_fused_smem_bytes(tactic, ab_stages), _EPI_ALIGN_BYTES) + cursor += tactic.mma_n * qtail.rope_dim * 4 + cursor = _align_up(cursor, _EPI_ALIGN_BYTES) + tactic.mma_n * 4 + if _qtail_pair(tactic, qtail): + cursor = _align_up(cursor, _EPI_ALIGN_BYTES) + tactic.mma_n * 4 + cursor = _align_up(cursor, _EPI_ALIGN_BYTES) + tactic.mma_n * 4 + cursor = _align_up(cursor, _MBARRIER_BYTES) + _MBARRIER_BYTES + return cursor + + +def idx_qtail_tactic(tactic: SplitKTactic, qtail: IndexerQTail) -> SplitKTactic: + """``tactic`` with at most as many A/B stages as fit beside the indexer Q tail's buffers.""" + if _qtail_pair(tactic, qtail) and 2 * tactic.split_k > 8: + raise ValueError(f"half-head tiles need a cluster of at most 8 CTAs: {tactic}") + stages = tactic.ab_stages + while stages > 1 and _qtail_smem_bytes(tactic, stages, qtail) > _SMEM_CAPACITY_BYTES: + stages -= 1 + return dataclasses.replace(tactic, ab_stages=stages) + + +#: A/B pipeline stages of the absorbed q_b queued beside the DSA indexer top-k. +_BESIDE_TOPK_AB_STAGES = 3 + +#: The launch plans of the fused q_b (``q_b_plan``). +Q_B_PLANS = ("regular", "beside_topk", "early_gather") + + +def q_b_plan(plan: str, m: int, n: int, k: int, device_index: int) -> SplitKTactic: + """The split-K tactic of the fused q_b ``[m, k] @ [k, n]`` under ``plan``. + + ``regular``: ``splitk_default_tactic``. ``beside_topk``: 128-row tiles, no split-K and three + A/B stages, so that a q_b queued beside the DSA indexer top-k takes few SMs and little memory + bandwidth from it; the regular tactic where that one does not serve the shape. + ``early_gather``: the source default (split-K 4) where the regular tactic is split-K 8, which + leaves the clusters of a flash decode that gathers before its grid dependency wait room to + land while q_b runs. ``run_splitk_fused`` then lowers the A/B stages to fit its buffers. + """ + if plan not in Q_B_PLANS: + raise ValueError(f"unknown q_b plan {plan!r}; expected one of {Q_B_PLANS}") + sm_count, smem_capacity = device_limits(device_index) + tactic = splitk_default_tactic(m, n, k, sm_count=sm_count, smem_capacity=smem_capacity) + if plan == "beside_topk": + candidate = dataclasses.replace( + tactic, mma_m=128, split_k=1, ab_stages=_BESIDE_TOPK_AB_STAGES + ) + try: + validate_tactic(candidate, m, n, k) + except ValueError: + return tactic + return candidate + if plan == "early_gather" and tactic.split_k == 8: + return default_tactic(m, n, k) + return tactic + + +def idx_qtail_plan(qtail: IndexerQTail, m: int, n: int, k: int, device_index: int) -> SplitKTactic: + """The split-K tactic of the indexer wq_b ``[m, k] @ [k, n]`` with the Q tail: the + ``splitk_default_tactic`` of the shape, its 64-row tiles as half-head tiles, its A/B stages + lowered to fit the Q tail's buffers. Raises ``ValueError`` where it does not serve the shape. + """ + sm_count, smem_capacity = device_limits(device_index) + tactic = splitk_default_tactic(m, n, k, sm_count=sm_count, smem_capacity=smem_capacity) + tactic = idx_qtail_tactic(tactic, qtail) + validate_tactic(tactic, m, n, k) + return tactic + + +# ----------------------------------------------------------------------------- PTX helpers + + +@dsl_user_op +def _ld_global_v4(addr: Int64, *, loc=None, ip=None): + """16-byte global load as four 32-bit words.""" + res = llvm.inline_asm( + llvm.StructType.get_literal([T.i32(), T.i32(), T.i32(), T.i32()]), + [addr.ir_value(loc=loc, ip=ip)], + "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + return [Int32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(4)] + + +@dsl_user_op +def _bf16x2_to_f32x2(word: Int32, *, loc=None, ip=None): + """The two BF16 halves of a 32-bit word (low half first) as FP32.""" + res = llvm.inline_asm( + llvm.StructType.get_literal([T.f32(), T.f32()]), + [word.ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 bflo, bfhi;\n\tshl.b32 bflo, $2, 16;\n\tand.b32 bfhi, $2, 0xFFFF0000;\n\t" + "mov.b32 $0, bflo;\n\tmov.b32 $1, bfhi;\n\t}", + "=f,=f,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + return [Float32(llvm.extractvalue(T.f32(), res, [i], loc=loc, ip=ip)) for i in range(2)] + + +@dsl_user_op +def _cvt_bf16x2(lo: Float32, hi: Float32, *, loc=None, ip=None) -> Int32: + """Round two FP32 values to BF16 (nearest even), ``lo`` in the low half.""" + return Int32( + llvm.inline_asm( + T.i32(), + [lo.ir_value(loc=loc, ip=ip), hi.ir_value(loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $2, $1;", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _cvt_e4m3x4(v0: Float32, v1: Float32, v2: Float32, v3: Float32, *, loc=None, ip=None) -> Int32: + """Four FP32 values to saturated E4M3 (nearest even), ``v0`` in the lowest byte.""" + return Int32( + llvm.inline_asm( + T.i32(), + [v.ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "{\n\t.reg .b16 fp8lo, fp8hi;\n\tcvt.rn.satfinite.e4m3x2.f32 fp8lo, $2, $1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 fp8hi, $4, $3;\n\tmov.b32 $0, {fp8lo, fp8hi};\n\t}", + "=r,f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _st_global_b32(addr: Int64, v: Int32, *, loc=None, ip=None) -> None: + llvm.inline_asm( + None, + [addr.ir_value(loc=loc, ip=ip), v.ir_value(loc=loc, ip=ip)], + "st.global.b32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _st_global_v2(addr: Int64, v0: Int32, v1: Int32, *, loc=None, ip=None) -> None: + llvm.inline_asm( + None, + [addr.ir_value(loc=loc, ip=ip), v0.ir_value(loc=loc, ip=ip), v1.ir_value(loc=loc, ip=ip)], + "st.global.v2.b32 [$0], {$1, $2};", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _st_global_v4( + addr: Int64, v0: Int32, v1: Int32, v2: Int32, v3: Int32, *, loc=None, ip=None +) -> None: + llvm.inline_asm( + None, + [addr.ir_value(loc=loc, ip=ip)] + [v.ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "st.global.v4.b32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _fma_rn(a: Float32, b: Float32, c: Float32, *, loc=None, ip=None) -> Float32: + """``a * b + c`` rounded once (nearest even).""" + return Float32( + llvm.inline_asm( + T.f32(), + [v.ir_value(loc=loc, ip=ip) for v in (a, b, c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _mul_rn(a: Float32, b: Float32, *, loc=None, ip=None) -> Float32: + """``a * b`` rounded once (nearest even), never contracted into an FMA.""" + return Float32( + llvm.inline_asm( + T.f32(), + [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _neg_f32(a: Float32, *, loc=None, ip=None) -> Float32: + """``-a`` (sign flip, exact).""" + return Float32( + llvm.inline_asm( + T.f32(), + [a.ir_value(loc=loc, ip=ip)], + "neg.f32 $0, $1;", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _fp8_ratio(amax: Float32, *, loc=None, ip=None) -> Float32: + """max(amax, 1e-12) * fp32(1 / 448): the per-row E4M3 ratio of the indexer prologue + (``_fp8_row_scale``), with its FP32 constants.""" + return Float32( + llvm.inline_asm( + T.f32(), + [amax.ir_value(loc=loc, ip=ip)], + "{\n\t.reg .f32 fpfloor;\n\tmax.f32 fpfloor, $1, 0f2B8CBCCC;\n\t" + "mul.rn.f32 $0, fpfloor, 0f3B124925;\n\t}", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _ue8m0_up(ratio: Float32, *, loc=None, ip=None) -> Float32: + """The positive ``ratio`` with its mantissa cleared and, if the mantissa was nonzero, its + exponent raised by one: the UE8M0 scale of the indexer prologue.""" + return Float32( + llvm.inline_asm( + T.f32(), + [ratio.ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 uebits, uemant, ueexp;\n\t.reg .pred uenz;\n\tmov.b32 uebits, $1;\n\t" + "and.b32 uemant, uebits, 0x007FFFFF;\n\tand.b32 ueexp, uebits, 0x7F800000;\n\t" + "setp.ne.u32 uenz, uemant, 0;\n\t@uenz add.u32 ueexp, ueexp, 0x00800000;\n\t" + "mov.b32 $0, ueexp;\n\t}", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +def _bf16_round(lo: Float32, hi: Float32): + """Two FP32 values rounded to BF16 (nearest even), back as FP32.""" + return _bf16x2_to_f32x2(_cvt_bf16x2(lo, hi)) + + +# --------------------------------------------------------------------------------- kernel + + +class SplitKFusedEpilogueGemmKernel(SplitKDenseGemmKernel): + """``SplitKDenseGemmKernel`` (no bias) with a fused owner epilogue.""" + + def __init__( + self, + *, + tactic: SplitKTactic, + use_pdl: bool, + epilogue: SplitKEpilogue, + reserve_full_smem: bool = True, + ) -> None: + super().__init__( + tactic=tactic, use_pdl=use_pdl, has_bias=False, reserve_full_smem=reserve_full_smem + ) + epilogue.validate() + self.epi = epilogue + self.launch_smem_bytes = min( + _align_up(_fused_smem_bytes(tactic, tactic.ab_stages), 1024) + + _LAUNCH_SMEM_HEADROOM_BYTES, + _SMEM_CAPACITY_BYTES, + ) + if (self.cta_m * self.cta_n // 4) % self.epilog_threads: + raise ValueError( + f"tile ({self.cta_m}, {self.cta_n}) does not split into 4-column groups" + ) + self.rows_per_tile = self.cta_n + self.threads_per_row = self.epilog_threads // self.cta_n + if epilogue.row_rms and ( + self.threads_per_row & (self.threads_per_row - 1) + or self.threads_per_row > 32 + or (epilogue.rms_k // _BF16_PER_VEC) % self.threads_per_row + ): + raise ValueError( + f"row_rms cannot split K={epilogue.rms_k} over a ({self.cta_m}, {self.cta_n}) tile" + ) + + @cute.experimental.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + gamma: cute.Tensor, + norm_out: cute.Tensor, + kv_lens: cute.Tensor, + cos_sin: cute.Tensor, + q_scale: cute.Tensor, + stream: _cuda.CUstream, + ): + self.kernel(a, b, c, gamma, norm_out, kv_lens, cos_sin, q_scale).launch( + grid=( + cute.ceil_div(c.layout.shape[0], self.cta_m), + cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k, + c.layout.shape[2], + ), + block=(self.threads_per_cta, 1, 1), + cluster=self.cluster_shape, + smem=cute.Int64( + utils.get_smem_capacity_in_bytes("sm_100") + if self.reserve_full_smem + else self.launch_smem_bytes + ), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.experimental.kernel + def kernel( + self, + mA: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L), K-major: the weight + mB: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L), K-major: the activation (rows = tokens) + mC: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L), M-major: BF16, or uint8 for fp8_rope + mGamma: cute.Tensor, # [K] BF16 norm weight (rms_out) + mNormOut: cute.Tensor, # [M, K] BF16 normalized activation (rms_out) + mKvLens: cute.Tensor, # [sequences] int32 KV lengths (fp8_rope) + mCosSin: cute.Tensor, # flat FP32 (cos, sin) table (fp8_rope) + mQScale: cute.Tensor, # [1] FP32 quantization scale (fp8_rope) + ): + """Allocate storage and dispatch the specialized warps (as the base kernel).""" + stages = self.num_ab_stage + + ab_dtype = mA.element_type + tiled_mma = sm100_utils.make_trivial_tiled_mma( + ab_dtype, + ab_dtype, + utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + + mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k) + block_idx = cute.arch.block_idx() + bidx = block_idx[0] + split_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + n_idx = block_idx[1] // self.split_k + l_idx = block_idx[2] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + sA = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + sB = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + + acc_layout = cute_ext.make_tmem_layout_acc(tiled_mma, self.mma_tiler_mn, acc_stage=1) + c_tiler_mn = (self.cta_m, self.cta_n) + + bar_full = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_empty = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tma_epilog = cute_ext.allocate( + cutlass.Int64, cute.AddressSpace.smem, cute.make_layout(1), alignment=_MBARRIER_BYTES + ).iterator + bar_mma_epilog = cute_ext.allocate( + cutlass.Int64, cute.AddressSpace.smem, cute.make_layout(1), alignment=_MBARRIER_BYTES + ).iterator + bar_tmem_alloc = cute_ext.allocate( + cutlass.Int64, cute.AddressSpace.smem, cute.make_layout(1), alignment=_MBARRIER_BYTES + ).iterator + tmem_base_ptr = cute_ext.allocate( + cutlass.Int32, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_TMEM_POINTER_BYTES, + ).iterator + + if cutlass.const_expr(self.split_k > 1): + mailbox = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.mailbox_elements), + alignment=_MAILBOX_ALIGN_BYTES, + ) + bar_reduce = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + else: + mailbox = sA + bar_reduce = bar_mma_epilog + + # Epilogue staging: the reduced FP32 tile (M-major, as the output tile) and the + # per-row scales. + sAcc = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout((self.cta_m, self.cta_n)), + alignment=_EPI_ALIGN_BYTES, + ) + sRstd = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.cta_n), + alignment=_EPI_ALIGN_BYTES, + ) + + if warp_idx == 0: + with cute.arch.elect_one(): + for i in range(stages): + cute.arch.mbarrier_init(bar_full + i, 2) + cute.arch.mbarrier_init(bar_empty + i, 1) + cute.arch.mbarrier_init(bar_tma_epilog, 32) + cute.arch.mbarrier_init(bar_mma_epilog, 1) + cute.arch.mbarrier_init(bar_tmem_alloc, 160) + + if cutlass.const_expr(self.split_k > 1): + cute.arch.mbarrier_init(bar_reduce, 1) + + cute.arch.mbarrier_init_fence() + if cutlass.const_expr(self.split_k > 1): + cute.arch.cluster_arrive_relaxed() + else: + cute.arch.barrier() + + k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k + k_tile_start = split_rank * k_tile_count + + if cutlass.const_expr(self.split_k > 1): + cute.arch.cluster_wait() + + if warp_idx == 0: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), + sA, + cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), + k_tile_start, + k_tile_count, + True, + ) + elif warp_idx == 1: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), + sB, + cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), + k_tile_start, + k_tile_count, + False, + ) + elif warp_idx == 2: + self.mma_warp( + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma, + sA, + sB, + tmem_base_ptr, + acc_layout, + self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), + k_tile_count, + ) + elif warp_idx >= 4: + self.epilog_fused( + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout, + cute.local_tile(mC, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.arch.thread_idx()[0] - 128, + utils.LayoutEnum.from_tensor(mC), + mailbox, + bar_reduce, + split_rank, + sAcc, + sRstd, + mB, + mC, + mGamma, + mNormOut, + mKvLens, + mCosSin, + mQScale, + bidx, + n_idx, + ) + + @cute.experimental.jit + def epilog_fused( + self, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + gD_tile: cute.Tensor, + epi_tid: cutlass.Int32, + d_layout: cutlass.Constexpr, + mailbox, + bar_reduce, + split_rank: cutlass.Int32, + sAcc: cute.Tensor, + sRstd: cute.Tensor, + mB: cute.Tensor, + mC: cute.Tensor, + mGamma: cute.Tensor, + mNormOut: cute.Tensor, + mKvLens: cute.Tensor, + mCosSin: cute.Tensor, + mQScale: cute.Tensor, + bidx: cutlass.Int32, + n_idx: cutlass.Int32, + ): + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) + + acc_view = cute.make_tensor( + cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr), + acc_layout, + )[((None, None), 0, 0, 0)] + + epi_tile = (self.cta_m, self.cta_n) + # The TMEM load op of the BF16 kernel, whatever the output type: the tile is staged + # in shared memory, so only the accumulator partition matters. + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( + sm100_utils.get_tmem_load_op( + (self.cta_m, self.cta_n, self.cta_k), + d_layout, + cutlass.BFloat16, + self.acc_dtype, + epi_tile, + False, + ), + acc_view, + ) + gD_epi = cute.flat_divide(gD_tile, epi_tile) + rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid) + rAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + thr_t2r = tiled_copy_t2r.get_slice(epi_tid) + + # Row statistics overlap the mainloop: only the owner needs them. + if cutlass.const_expr(self.epi.row_rms): + if split_rank == OWNER_RANK: + self._row_rstd(mB, sRstd, n_idx, epi_tid) + if cutlass.const_expr(self.epi.rms_out): + self._rms_out(mB, mC, mGamma, mNormOut, sRstd, bidx, n_idx, epi_tid) + + cute.arch.mbarrier_wait(bar_mma_epilog, 0) + cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) + cute.arch.fence_view_async_tmem_load() + cute.arch.mbarrier_arrive(bar_tmem_alloc) + + if cutlass.const_expr(self.split_k > 1): + values_per_thread = cutlass.const_expr(cute.size(rmem_layout)) + values_per_peer = cutlass.const_expr(self.epilog_threads * values_per_thread) + if split_rank != OWNER_RANK: + for value_idx in cutlass.range_constexpr(0, values_per_thread, 4): + _store_shared_remote_v4( + rAcc[value_idx], + rAcc[value_idx + 1], + rAcc[value_idx + 2], + rAcc[value_idx + 3], + mailbox.iterator + + (split_rank - Int32(1)) * values_per_peer + + epi_tid * values_per_thread + + value_idx, + bar_reduce, + Int32(OWNER_RANK), + ) + else: + if epi_tid == 0: + cute.arch.mbarrier_arrive_and_expect_tx( + bar_reduce, self.expected_transaction_bytes + ) + cute.arch.mbarrier_wait(bar_reduce, 0) + for peer in cutlass.range_constexpr(self.split_k - 1): + for value_idx in cutlass.range_constexpr(values_per_thread): + rAcc[value_idx] = ( + rAcc[value_idx] + + mailbox[ + peer * values_per_peer + epi_tid * values_per_thread + value_idx + ] + ) + + if split_rank == OWNER_RANK: + # Stage the tile by output coordinate, then store four columns per thread. + cute_ext.partition_and_copy(thr_t2r, rAcc, sAcc) + cute.arch.barrier(barrier_id=_EPI_BARRIER_ID, number_of_threads=self.epilog_threads) + self._store_tile(sAcc, sRstd, mC, mKvLens, mCosSin, mQScale, bidx, n_idx, epi_tid) + + @cute.experimental.jit + def _row_rstd( + self, + mB: cute.Tensor, + sRstd: cute.Tensor, + n_idx: cutlass.Int32, + epi_tid: cutlass.Int32, + ): + """sRstd[row] = rsqrt(mean(x[row]^2) + eps) for the tile's rows; ends with a barrier.""" + tpr = cutlass.const_expr(self.threads_per_row) + vecs_per_thread = cutlass.const_expr(self.epi.rms_k // _BF16_PER_VEC // tpr) + # The activation is the previous grid's output. + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + row = epi_tid // tpr + lane = epi_tid - row * tpr + token = n_idx * self.cta_n + row + ssq = Float32(0.0) + if token < cute.size(mB, mode=[0]): + base = (mB.iterator + cute.crd2idx((token, 0, 0), mB.layout)).toint() + Int64(lane * 16) + for i in cutlass.range_constexpr(vecs_per_thread): + words = _ld_global_v4(base + Int64(i * tpr * 16)) + for j in cutlass.range_constexpr(4): + pair = _bf16x2_to_f32x2(words[j]) + ssq = ssq + pair[0] * pair[0] + pair[1] * pair[1] + for i in cutlass.range_constexpr(int(math.log2(tpr))): + ssq = ssq + cute.arch.shuffle_sync_bfly(ssq, offset=1 << i) + if lane == 0: + sRstd[row] = cute.math.rsqrt( + ssq / Float32(self.epi.rms_k) + Float32(self.epi.rms_eps), fastmath=True + ) + cute.arch.barrier(barrier_id=_EPI_BARRIER_ID, number_of_threads=self.epilog_threads) + + @cute.experimental.jit + def _rms_out( + self, + mB: cute.Tensor, + mC: cute.Tensor, + mGamma: cute.Tensor, + mNormOut: cute.Tensor, + sRstd: cute.Tensor, + bidx: cutlass.Int32, + n_idx: cutlass.Int32, + epi_tid: cutlass.Int32, + ): + """Write x * rstd * gamma for this tile's rows and its share of the K columns.""" + vecs = cutlass.const_expr(self.epi.rms_k // _BF16_PER_VEC) + n_tiles = cute.ceil_div(cute.size(mC, mode=[0]), self.cta_m) + chunk = (vecs + n_tiles - 1) // n_tiles + total = chunk * self.cta_n + iters = (total + self.epilog_threads - 1) // self.epilog_threads + tokens = cute.size(mB, mode=[0]) + for i in cutlass.range(iters, unroll=1): + idx = epi_tid + i * self.epilog_threads + row = idx // chunk + vec = bidx * chunk + (idx - row * chunk) + token = n_idx * self.cta_n + row + if idx < total: + if vec < vecs: + if token < tokens: + col = vec * _BF16_PER_VEC + x_words = _ld_global_v4( + (mB.iterator + cute.crd2idx((token, col, 0), mB.layout)).toint() + ) + g_words = _ld_global_v4((mGamma.iterator + col).toint()) + rstd = sRstd[row] + out_words = [] + for j in cutlass.range_constexpr(4): + x_pair = _bf16x2_to_f32x2(x_words[j]) + g_pair = _bf16x2_to_f32x2(g_words[j]) + out_words.append( + _cvt_bf16x2( + x_pair[0] * rstd * g_pair[0], x_pair[1] * rstd * g_pair[1] + ) + ) + out_ptr = mNormOut.iterator + cute.crd2idx((token, col), mNormOut.layout) + _st_global_v4( + out_ptr.toint(), + out_words[0], + out_words[1], + out_words[2], + out_words[3], + ) + + @cute.experimental.jit + def _store_tile( + self, + sAcc: cute.Tensor, + sRstd: cute.Tensor, + mC: cute.Tensor, + mKvLens: cute.Tensor, + mCosSin: cute.Tensor, + mQScale: cute.Tensor, + bidx: cutlass.Int32, + n_idx: cutlass.Int32, + epi_tid: cutlass.Int32, + ): + """Store the staged tile: four consecutive output columns of one row per thread.""" + groups_m = cutlass.const_expr(self.cta_m // 4) + groups_per_thread = cutlass.const_expr(self.cta_m * self.cta_n // 4 // self.epilog_threads) + n_feat = cute.size(mC, mode=[0]) + n_tok = cute.size(mC, mode=[1]) + q_scale = Float32(1.0) + if cutlass.const_expr(self.epi.fp8_rope): + q_scale = mQScale[0] + for g in cutlass.range_constexpr(groups_per_thread): + gi = epi_tid + g * self.epilog_threads + row = gi // groups_m + m4 = (gi - row * groups_m) * 4 + v0 = sAcc[m4, row] + v1 = sAcc[m4 + 1, row] + v2 = sAcc[m4 + 2, row] + v3 = sAcc[m4 + 3, row] + if cutlass.const_expr(self.epi.row_rms): + rstd = sRstd[row] + v0 = v0 * rstd + v1 = v1 * rstd + v2 = v2 * rstd + v3 = v3 * rstd + token = n_idx * self.cta_n + row + feature = bidx * self.cta_m + m4 + if token < n_tok: + if feature < n_feat: + addr = (mC.iterator + cute.crd2idx((feature, token, 0), mC.layout)).toint() + if cutlass.const_expr(self.epi.fp8_rope): + nope = cutlass.const_expr(self.epi.head_dim - self.epi.rope_dim) + col = feature % self.epi.head_dim + if col >= nope: + seq = token // self.epi.seq_len + pos = mKvLens[seq] - self.epi.seq_len + (token - seq * self.epi.seq_len) + # The table holds rope_dim (cos, sin) pairs per position; pair p + # rotates columns 2p and 2p + 1 of the head's RoPE slice. + cs = (pos * self.epi.rope_dim + (col - nope) // 2) * 2 + c0 = mCosSin[cs] + s0 = mCosSin[cs + 1] + c1 = mCosSin[cs + 2] + s1 = mCosSin[cs + 3] + r0 = c0 * v0 - s0 * v1 + r1 = c0 * v1 + s0 * v0 + r2 = c1 * v2 - s1 * v3 + r3 = c1 * v3 + s1 * v2 + v0 = r0 + v1 = r1 + v2 = r2 + v3 = r3 + _st_global_b32( + addr, + _cvt_e4m3x4(v0 * q_scale, v1 * q_scale, v2 * q_scale, v3 * q_scale), + ) + else: + _st_global_v2(addr, _cvt_bf16x2(v0, v1), _cvt_bf16x2(v2, v3)) + + +class SplitKIndexerQTailGemmKernel(SplitKFusedEpilogueGemmKernel): + """The ``row_rms`` kernel of the DSA indexer wq_b with the indexer Q tail (``IndexerQTail``) + in the owner epilogue, which writes the E4M3 Q, its per-head scale and the scaled attention + weights instead of the BF16 output. + + A tile is one head (``mma_m == head_dim``: a warp finishes one token's head) or half a head + (``2 * mma_m == head_dim``: half a warp per token, and the two halves of a head run in one + cluster of shape (2, split_k, 1), whose owners exchange their row amax through DSMEM). The + mainloop, the split-K reduction and the row RMS are those of the fused kernel, so the GEMM of + a tactic is unchanged. + """ + + def __init__( + self, + *, + tactic: SplitKTactic, + use_pdl: bool, + epilogue: SplitKEpilogue, + qtail: IndexerQTail, + reserve_full_smem: bool = True, + ) -> None: + super().__init__( + tactic=tactic, use_pdl=use_pdl, epilogue=epilogue, reserve_full_smem=reserve_full_smem + ) + qtail.validate() + if not epilogue.row_rms or epilogue.rms_out or epilogue.fp8_rope: + raise ValueError("the indexer Q tail follows the row_rms epilogue alone") + self.qtail = qtail + self.pair = _qtail_pair(tactic, qtail) + self.cluster_shape = (2 if self.pair else 1, tactic.split_k, 1) + self.lanes_per_row = self.cta_m // 4 + if _WARP_LANES % self.lanes_per_row or (self.epilog_threads // self.cta_n) == 0: + raise ValueError(f"tile ({self.cta_m}, {self.cta_n}) does not map rows onto lanes") + if qtail.rope_dim % (self.epilog_threads // self.cta_n): + raise ValueError(f"rope_dim {qtail.rope_dim} does not split over a tile row's threads") + self.launch_smem_bytes = min( + _align_up(_qtail_smem_bytes(tactic, tactic.ab_stages, qtail), 1024) + + _LAUNCH_SMEM_HEADROOM_BYTES, + _SMEM_CAPACITY_BYTES, + ) + + @cute.experimental.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + positions: cute.Tensor, + cos_sin: cute.Tensor, + weights_in: cute.Tensor, + q_scale_out: cute.Tensor, + weights_out: cute.Tensor, + stream: _cuda.CUstream, + ): + self.kernel(a, b, c, positions, cos_sin, weights_in, q_scale_out, weights_out).launch( + grid=( + cute.ceil_div(c.layout.shape[0], self.cta_m), + cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k, + c.layout.shape[2], + ), + block=(self.threads_per_cta, 1, 1), + cluster=self.cluster_shape, + smem=cute.Int64( + utils.get_smem_capacity_in_bytes("sm_100") + if self.reserve_full_smem + else self.launch_smem_bytes + ), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.experimental.kernel + def kernel( + self, + mA: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L), K-major: the weight + mB: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L), K-major: the activation (rows = tokens) + mC: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L), M-major: uint8 (E4M3 bits) + mPos: cute.Tensor, # [tokens] int32 positions + mCosSin: cute.Tensor, # [positions, >= rope_dim] FP32 [cos | sin] rows + mWIn: cute.Tensor, # [tokens, heads] BF16 attention weights (weights_proj output) + mQScaleOut: cute.Tensor, # [tokens, heads] FP32 E4M3 scales + mWOut: cute.Tensor, # [tokens, heads] FP32 attention weights * scale * weight_scale + ): + """Allocate storage and dispatch the specialized warps (as the fused kernel).""" + stages = self.num_ab_stage + + ab_dtype = mA.element_type + tiled_mma = sm100_utils.make_trivial_tiled_mma( + ab_dtype, + ab_dtype, + utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + + mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k) + block_idx = cute.arch.block_idx() + bidx = block_idx[0] + # The cluster's ranks: x is the half of the head (half-head tiles), y the K slice. + split_rank = cute.arch.make_warp_uniform(block_idx[1] % self.split_k) + half = Int32(0) + if cutlass.const_expr(self.pair): + half = cute.arch.make_warp_uniform(bidx % 2) + n_idx = block_idx[1] // self.split_k + l_idx = block_idx[2] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + sA = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + sB = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + + acc_layout = cute_ext.make_tmem_layout_acc(tiled_mma, self.mma_tiler_mn, acc_stage=1) + c_tiler_mn = (self.cta_m, self.cta_n) + + bar_full = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_empty = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tma_epilog = cute_ext.allocate( + cutlass.Int64, cute.AddressSpace.smem, cute.make_layout(1), alignment=_MBARRIER_BYTES + ).iterator + bar_mma_epilog = cute_ext.allocate( + cutlass.Int64, cute.AddressSpace.smem, cute.make_layout(1), alignment=_MBARRIER_BYTES + ).iterator + bar_tmem_alloc = cute_ext.allocate( + cutlass.Int64, cute.AddressSpace.smem, cute.make_layout(1), alignment=_MBARRIER_BYTES + ).iterator + tmem_base_ptr = cute_ext.allocate( + cutlass.Int32, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_TMEM_POINTER_BYTES, + ).iterator + + if cutlass.const_expr(self.split_k > 1): + mailbox = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.mailbox_elements), + alignment=_MAILBOX_ALIGN_BYTES, + ) + bar_reduce = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + else: + mailbox = sA + bar_reduce = bar_mma_epilog + + sAcc = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout((self.cta_m, self.cta_n)), + alignment=_EPI_ALIGN_BYTES, + ) + sRstd = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.cta_n), + alignment=_EPI_ALIGN_BYTES, + ) + # The tile's tokens: their cos/sin rows and attention weights, loaded while the mainloop + # runs. + sCS = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.cta_n * self.qtail.rope_dim), + alignment=_EPI_ALIGN_BYTES, + ) + sW = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.cta_n), + alignment=_EPI_ALIGN_BYTES, + ) + if cutlass.const_expr(self.pair): + # Row amax of this half, and of the other half (written by its owner). + sAmaxMine = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.cta_n), + alignment=_EPI_ALIGN_BYTES, + ) + sAmaxPeer = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.cta_n), + alignment=_EPI_ALIGN_BYTES, + ) + bar_amax = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + else: + sAmaxMine = sRstd + sAmaxPeer = sRstd + bar_amax = bar_mma_epilog + + if warp_idx == 0: + with cute.arch.elect_one(): + for i in range(stages): + cute.arch.mbarrier_init(bar_full + i, 2) + cute.arch.mbarrier_init(bar_empty + i, 1) + cute.arch.mbarrier_init(bar_tma_epilog, 32) + cute.arch.mbarrier_init(bar_mma_epilog, 1) + cute.arch.mbarrier_init(bar_tmem_alloc, 160) + + if cutlass.const_expr(self.split_k > 1): + cute.arch.mbarrier_init(bar_reduce, 1) + if cutlass.const_expr(self.pair): + # The local arrival plus the other owner's transaction bytes. + cute.arch.mbarrier_init(bar_amax, 1) + + cute.arch.mbarrier_init_fence() + if cutlass.const_expr(self.split_k > 1 or self.pair): + cute.arch.cluster_arrive_relaxed() + else: + cute.arch.barrier() + + k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k + k_tile_start = split_rank * k_tile_count + + if cutlass.const_expr(self.split_k > 1 or self.pair): + cute.arch.cluster_wait() + + if warp_idx == 0: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), + sA, + cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), + k_tile_start, + k_tile_count, + True, + ) + elif warp_idx == 1: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), + sB, + cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), + k_tile_start, + k_tile_count, + False, + ) + elif warp_idx == 2: + self.mma_warp( + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma, + sA, + sB, + tmem_base_ptr, + acc_layout, + self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), + k_tile_count, + ) + elif warp_idx >= 4: + self.epilog_qtail( + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout, + cute.local_tile(mC, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.arch.thread_idx()[0] - 128, + utils.LayoutEnum.from_tensor(mC), + mailbox, + bar_reduce, + split_rank, + half, + sAcc, + sRstd, + sCS, + sW, + sAmaxMine, + sAmaxPeer, + bar_amax, + mB, + mC, + mPos, + mCosSin, + mWIn, + mQScaleOut, + mWOut, + bidx, + n_idx, + ) + + @cute.experimental.jit + def epilog_qtail( + self, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + gD_tile: cute.Tensor, + epi_tid: cutlass.Int32, + d_layout: cutlass.Constexpr, + mailbox, + bar_reduce, + split_rank: cutlass.Int32, + half: cutlass.Int32, + sAcc: cute.Tensor, + sRstd: cute.Tensor, + sCS: cute.Tensor, + sW: cute.Tensor, + sAmaxMine: cute.Tensor, + sAmaxPeer: cute.Tensor, + bar_amax, + mB: cute.Tensor, + mC: cute.Tensor, + mPos: cute.Tensor, + mCosSin: cute.Tensor, + mWIn: cute.Tensor, + mQScaleOut: cute.Tensor, + mWOut: cute.Tensor, + bidx: cutlass.Int32, + n_idx: cutlass.Int32, + ): + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) + + acc_view = cute.make_tensor( + cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr), + acc_layout, + )[((None, None), 0, 0, 0)] + + epi_tile = (self.cta_m, self.cta_n) + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( + sm100_utils.get_tmem_load_op( + (self.cta_m, self.cta_n, self.cta_k), + d_layout, + cutlass.BFloat16, + self.acc_dtype, + epi_tile, + False, + ), + acc_view, + ) + gD_epi = cute.flat_divide(gD_tile, epi_tile) + rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid) + rAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + thr_t2r = tiled_copy_t2r.get_slice(epi_tid) + + head = bidx + if cutlass.const_expr(self.pair): + head = bidx // 2 + # The owner's row statistics, cos/sin rows and attention weights overlap the mainloop. + if split_rank == OWNER_RANK: + self._row_rstd(mB, sRstd, n_idx, epi_tid) + self._qtail_prefetch(mB, mPos, mCosSin, mWIn, sCS, sW, half, head, n_idx, epi_tid) + + cute.arch.mbarrier_wait(bar_mma_epilog, 0) + cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) + cute.arch.fence_view_async_tmem_load() + cute.arch.mbarrier_arrive(bar_tmem_alloc) + + if cutlass.const_expr(self.split_k > 1): + # The owner of this tile is the cluster's CTA of the same half at K slice 0. + owner_rank = half + values_per_thread = cutlass.const_expr(cute.size(rmem_layout)) + values_per_peer = cutlass.const_expr(self.epilog_threads * values_per_thread) + if split_rank != OWNER_RANK: + for value_idx in cutlass.range_constexpr(0, values_per_thread, 4): + _store_shared_remote_v4( + rAcc[value_idx], + rAcc[value_idx + 1], + rAcc[value_idx + 2], + rAcc[value_idx + 3], + mailbox.iterator + + (split_rank - Int32(1)) * values_per_peer + + epi_tid * values_per_thread + + value_idx, + bar_reduce, + owner_rank, + ) + else: + if epi_tid == 0: + cute.arch.mbarrier_arrive_and_expect_tx( + bar_reduce, self.expected_transaction_bytes + ) + cute.arch.mbarrier_wait(bar_reduce, 0) + for peer in cutlass.range_constexpr(self.split_k - 1): + for value_idx in cutlass.range_constexpr(values_per_thread): + rAcc[value_idx] = ( + rAcc[value_idx] + + mailbox[ + peer * values_per_peer + epi_tid * values_per_thread + value_idx + ] + ) + + if split_rank == OWNER_RANK: + cute_ext.partition_and_copy(thr_t2r, rAcc, sAcc) + cute.arch.barrier(barrier_id=_EPI_BARRIER_ID, number_of_threads=self.epilog_threads) + self._store_qtail( + sAcc, + sRstd, + sCS, + sW, + sAmaxMine, + sAmaxPeer, + bar_amax, + mC, + mQScaleOut, + mWOut, + bidx, + half, + head, + n_idx, + epi_tid, + ) + + @cute.experimental.jit + def _qtail_prefetch( + self, + mB: cute.Tensor, + mPos: cute.Tensor, + mCosSin: cute.Tensor, + mWIn: cute.Tensor, + sCS: cute.Tensor, + sW: cute.Tensor, + half: cutlass.Int32, + head: cutlass.Int32, + n_idx: cutlass.Int32, + epi_tid: cutlass.Int32, + ): + """sCS[row] = the cos/sin row of the token's position and sW[row] = its attention weight, + for the tile's tokens, in the tile that holds the head's RoPE columns. Runs after + ``_row_rstd``, whose grid dependency wait covers the positions and the weights.""" + rope = cutlass.const_expr(self.qtail.rope_dim) + tpr = cutlass.const_expr(self.epilog_threads // self.cta_n) + per = cutlass.const_expr(rope // tpr) + row = epi_tid // tpr + lane = epi_tid - row * tpr + token = n_idx * self.cta_n + row + if token < cute.size(mB, mode=[0]): + if half == 0: + pos = mPos[token] + for i in cutlass.range_constexpr(per): + col = lane * per + i + sCS[row * rope + col] = mCosSin[pos, col] + if lane == 0: + sW[row] = mWIn[token, head].to(Float32) + + @cute.experimental.jit + def _store_qtail( + self, + sAcc: cute.Tensor, + sRstd: cute.Tensor, + sCS: cute.Tensor, + sW: cute.Tensor, + sAmaxMine: cute.Tensor, + sAmaxPeer: cute.Tensor, + bar_amax, + mC: cute.Tensor, + mQScaleOut: cute.Tensor, + mWOut: cute.Tensor, + bidx: cutlass.Int32, + half: cutlass.Int32, + head: cutlass.Int32, + n_idx: cutlass.Int32, + epi_tid: cutlass.Int32, + ): + """Finish the staged tile: four consecutive columns of one token row per thread, the + lanes of a row reducing the head's amax with butterfly shuffles.""" + qt = self.qtail + rope = cutlass.const_expr(qt.rope_dim) + rhalf = cutlass.const_expr(qt.rope_dim // 2) + lanes = cutlass.const_expr(self.lanes_per_row) + groups_per_thread = cutlass.const_expr(self.cta_m * self.cta_n // 4 // self.epilog_threads) + n_tok = cute.size(mC, mode=[1]) + col0 = half * self.cta_m + rows = [] + m4s = [] + vals = [] + amaxes = [] + for g in cutlass.range_constexpr(groups_per_thread): + gi = epi_tid + g * self.epilog_threads + row = gi // lanes + m4 = (gi - row * lanes) * 4 + rstd = sRstd[row] + # The BF16 wq_b output, as the BF16 store rounds it. + x01 = _bf16_round(sAcc[m4, row] * rstd, sAcc[m4 + 1, row] * rstd) + x23 = _bf16_round(sAcc[m4 + 2, row] * rstd, sAcc[m4 + 3, row] * rstd) + q0 = x01[0] + q1 = x01[1] + q2 = x23[0] + q3 = x23[1] + col = col0 + m4 + if col < rope: + # x * cos + (+-partner) * sin as fma(x, cos, rn(+-partner * sin)), then BF16. + if cutlass.const_expr(qt.interleave): + cs = row * rope + col // 2 + c0 = sCS[cs] + s0 = sCS[cs + rhalf] + c1 = sCS[cs + 1] + s1 = sCS[cs + rhalf + 1] + r01 = _bf16_round( + _fma_rn(q0, c0, _mul_rn(_neg_f32(q1), s0)), + _fma_rn(q1, c0, _mul_rn(q0, s0)), + ) + r23 = _bf16_round( + _fma_rn(q2, c1, _mul_rn(_neg_f32(q3), s1)), + _fma_rn(q3, c1, _mul_rn(q2, s1)), + ) + else: + # NeoX: column d pairs with d +- rope / 2; four columns share a side. + lower = col < rhalf + pm4 = m4 - rhalf + cs = row * rope + col - rhalf + if lower: + pm4 = m4 + rhalf + cs = row * rope + col + p01 = _bf16_round(sAcc[pm4, row] * rstd, sAcc[pm4 + 1, row] * rstd) + p23 = _bf16_round(sAcc[pm4 + 2, row] * rstd, sAcc[pm4 + 3, row] * rstd) + p0 = p01[0] + p1 = p01[1] + p2 = p23[0] + p3 = p23[1] + if lower: + p0 = _neg_f32(p0) + p1 = _neg_f32(p1) + p2 = _neg_f32(p2) + p3 = _neg_f32(p3) + r01 = _bf16_round( + _fma_rn(q0, sCS[cs], _mul_rn(p0, sCS[cs + rhalf])), + _fma_rn(q1, sCS[cs + 1], _mul_rn(p1, sCS[cs + rhalf + 1])), + ) + r23 = _bf16_round( + _fma_rn(q2, sCS[cs + 2], _mul_rn(p2, sCS[cs + rhalf + 2])), + _fma_rn(q3, sCS[cs + 3], _mul_rn(p3, sCS[cs + rhalf + 3])), + ) + q0 = r01[0] + q1 = r01[1] + q2 = r23[0] + q3 = r23[1] + amax = cute.arch.fmax( + cute.arch.fmax(abs(q0), abs(q1)), cute.arch.fmax(abs(q2), abs(q3)) + ) + for i in cutlass.range_constexpr(int(math.log2(lanes))): + amax = cute.arch.fmax( + amax, cute.arch.shuffle_sync_bfly(amax, offset=lanes >> (i + 1)) + ) + rows.append(row) + m4s.append(m4) + vals.append((q0, q1, q2, q3)) + amaxes.append(amax) + + if cutlass.const_expr(self.pair): + # The other half's row amax: each owner sends its rows' amax to the other owner. + for g in cutlass.range_constexpr(groups_per_thread): + if m4s[g] == 0: + sAmaxMine[rows[g]] = amaxes[g] + cute.arch.barrier(barrier_id=_EPI_BARRIER_ID, number_of_threads=self.epilog_threads) + if epi_tid < self.cta_n // 4: + v = epi_tid * 4 + _store_shared_remote_v4( + sAmaxMine[v], + sAmaxMine[v + 1], + sAmaxMine[v + 2], + sAmaxMine[v + 3], + sAmaxPeer.iterator + v, + bar_amax, + Int32(1) - half, + ) + if epi_tid == 0: + cute.arch.mbarrier_arrive_and_expect_tx(bar_amax, self.cta_n * 4) + cute.arch.mbarrier_wait(bar_amax, 0) + for g in cutlass.range_constexpr(groups_per_thread): + amaxes[g] = cute.arch.fmax(amaxes[g], sAmaxPeer[rows[g]]) + + for g in cutlass.range_constexpr(groups_per_thread): + row = rows[g] + m4 = m4s[g] + q0, q1, q2, q3 = vals[g] + token = n_idx * self.cta_n + row + if token < n_tok: + scale = _fp8_ratio(amaxes[g]) + if cutlass.const_expr(qt.ue8m0): + scale = _ue8m0_up(scale) + inv = cute.arch.rcp_approx(scale) + feature = bidx * self.cta_m + m4 + addr = (mC.iterator + cute.crd2idx((feature, token, 0), mC.layout)).toint() + _st_global_b32( + addr, + _cvt_e4m3x4( + _mul_rn(q0, inv), _mul_rn(q1, inv), _mul_rn(q2, inv), _mul_rn(q3, inv) + ), + ) + if col0 + m4 == 0: + mQScaleOut[token, head] = scale + mWOut[token, head] = _mul_rn(_mul_rn(sW[row], scale), Float32(qt.weight_scale)) + + +# ----------------------------------------------------------------------------------- host + + +@cute.experimental.jit +def _bmm_fused( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + gamma: cute.Tensor, + norm_out: cute.Tensor, + kv_lens: cute.Tensor, + cos_sin: cute.Tensor, + q_scale: cute.Tensor, + stream: _cuda.CUstream, +): + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])), + gamma, + norm_out, + kv_lens, + cos_sin, + q_scale, + stream, + ) + + +@functools.cache +def _placeholders(device_index: int) -> dict: + """Stand-ins for the operands an epilogue does not read (one compiled signature).""" + device = _torch.device("cuda", device_index) + return { + "gamma": _torch.zeros(_BF16_PER_VEC, dtype=_torch.bfloat16, device=device), + "norm_out": _torch.zeros((1, _BF16_PER_VEC), dtype=_torch.bfloat16, device=device), + "kv_lens": _torch.zeros(4, dtype=_torch.int32, device=device), + "cos_sin": _torch.zeros(4, dtype=_torch.float32, device=device), + "q_scale": _torch.ones(4, dtype=_torch.float32, device=device), + } + + +def _aux_cute_tensors(gamma, norm_out, kv_lens, cos_sin, q_scale): + # detach(): a module parameter (the norm weight) cannot be exported through DLPack otherwise. + return ( + _from_dlpack_dynamic(gamma.detach(), 0, assumed_align=16), + _from_dlpack_dynamic(norm_out.detach(), 1, assumed_align=16), + _from_dlpack_dynamic(kv_lens.detach(), 0, assumed_align=4), + _from_dlpack_dynamic(cos_sin.detach(), 0, assumed_align=4), + _from_dlpack_dynamic(q_scale.detach(), 0, assumed_align=4), + ) + + +def _swap_cute_tensors(a, b, out): + a_swap = b.detach().unsqueeze(0).transpose(-2, -1) + b_swap = a.detach().unsqueeze(0).transpose(-2, -1) + c_swap = out.detach().unsqueeze(0).transpose(-2, -1) + leading_dims = tuple(_detect_leading_dim(t) for t in (a_swap, b_swap, c_swap)) + tensors = tuple( + _from_dlpack_dynamic(t, ld) + for t, ld in zip((a_swap, b_swap, c_swap), leading_dims, strict=True) + ) + return tensors, leading_dims + + +@functools.cache +def _get_compiled_fused_kernel( + tactic: SplitKTactic, + use_pdl: bool, + epilogue: SplitKEpilogue, + leading_dims: tuple[int, int, int], + out_dtype: _torch.dtype, + reserve_full_smem: bool, + device_index: int, +): + kernel = SplitKFusedEpilogueGemmKernel( + tactic=tactic, use_pdl=use_pdl, epilogue=epilogue, reserve_full_smem=reserve_full_smem + ) + device = _torch.device("cuda", device_index) + m, n, k = 8, 64 * 9, _CTA_K * 2 + weight = _torch.empty((n, k), dtype=_torch.bfloat16, device=device) + act = _torch.empty((m, k), dtype=_torch.bfloat16, device=device) + out = _torch.empty((m, n), dtype=out_dtype, device=device) + tensors, repr_leading = _swap_cute_tensors(act, weight.t(), out) + if repr_leading != leading_dims: + raise ValueError(f"unsupported operand layouts {leading_dims} (expected {repr_leading})") + ph = _placeholders(device_index) + aux = _aux_cute_tensors( + ph["gamma"], ph["norm_out"], ph["kv_lens"], ph["cos_sin"], ph["q_scale"] + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(device).cuda_stream) + return cute_ext.compile(_bmm_fused, kernel, *tensors, *aux, stream) + + +def _check(cond: bool, msg: str) -> None: + if not cond: + raise ValueError(msg) + + +def _validate_fused(a, b, out, epilogue, gamma, norm_out, kv_lens, cos_sin, q_scale): + _check(all(isinstance(t, _torch.Tensor) for t in (a, b, out)), "a, b and out must be tensors") + _check(a.ndim == 2 and b.ndim == 2 and out.ndim == 2, "a, b and out must be 2D") + _check( + a.is_cuda and b.device == a.device and out.device == a.device, + "a, b and out must be on one CUDA device", + ) + _check(a.dtype == _torch.bfloat16 and b.dtype == _torch.bfloat16, "a and b must be BF16") + m, k = a.shape + _check(b.shape[0] == k, f"incompatible shapes {tuple(a.shape)} x {tuple(b.shape)}") + n = b.shape[1] + _check(tuple(out.shape) == (m, n), f"out must be {(m, n)}, got {tuple(out.shape)}") + _check(1 <= m <= 32 and k % _CTA_K == 0 and n % 4 == 0, f"unsupported shape M={m} N={n} K={k}") + # a may be a row view of a wider buffer (the raw q_a slice of the fused-A output). + _check( + a.stride(1) == 1 and a.stride(0) >= k and (a.stride(0) * 2) % 16 == 0, + f"a needs unit inner stride and a 16-byte row pitch, got strides {a.stride()}", + ) + _check(b.t().is_contiguous(), "b must be column-major (the transpose of a contiguous [N, K])") + _check(out.is_contiguous(), "out must be contiguous") + _check(all(t.data_ptr() % 32 == 0 for t in (a, b, out)), "a, b and out must be 32-byte aligned") + if epilogue.fp8_rope: + _check(out.dtype == _torch.uint8, "fp8_rope writes uint8 (E4M3 bits)") + _check( + n % epilogue.head_dim == 0, f"N={n} is not a multiple of head_dim {epilogue.head_dim}" + ) + _check(m % epilogue.seq_len == 0, f"M={m} is not a multiple of seq_len {epilogue.seq_len}") + for t, name, dtype in ( + (kv_lens, "kv_lens", _torch.int32), + (cos_sin, "cos_sin", _torch.float32), + (q_scale, "q_scale", _torch.float32), + ): + _check( + isinstance(t, _torch.Tensor) + and t.dtype == dtype + and t.is_contiguous() + and t.device == a.device, + f"fp8_rope needs a contiguous {dtype} {name} on the GEMM's device", + ) + _check(kv_lens.numel() >= m // epilogue.seq_len, "kv_lens has fewer entries than sequences") + _check(q_scale.numel() >= 1, "q_scale is empty") + else: + _check(out.dtype == _torch.bfloat16, "out must be BF16 without fp8_rope") + if epilogue.row_rms: + _check(k == epilogue.rms_k, f"row_rms was specialized for K={epilogue.rms_k}, got {k}") + if epilogue.rms_out: + _check( + isinstance(gamma, _torch.Tensor) + and gamma.dtype == _torch.bfloat16 + and gamma.is_contiguous() + and gamma.numel() == k + and gamma.data_ptr() % 16 == 0, + "rms_out needs a contiguous 16-byte aligned BF16 gamma of K elements", + ) + _check( + isinstance(norm_out, _torch.Tensor) + and norm_out.dtype == _torch.bfloat16 + and norm_out.is_contiguous() + and tuple(norm_out.shape) == (m, k) + and norm_out.data_ptr() % 16 == 0, + "rms_out needs a contiguous 16-byte aligned BF16 norm_out [M, K]", + ) + return m, n, k + + +def run_splitk_fused( + a: _torch.Tensor, + b: _torch.Tensor, + out: _torch.Tensor, + pdl: bool, + tactic: SplitKTactic, + epilogue: SplitKEpilogue, + *, + gamma: _torch.Tensor | None = None, + norm_out: _torch.Tensor | None = None, + kv_lens: _torch.Tensor | None = None, + cos_sin: _torch.Tensor | None = None, + q_scale: _torch.Tensor | None = None, + reserve_full_smem: bool = True, +) -> _torch.Tensor: + """Run ``a[M, K] @ b[K, N]`` with the split-K kernel and a fused epilogue into ``out``. + + Args: + a: Activation [M, K] BF16 with unit inner stride (may be a row view of a wider + buffer). With ``epilogue.row_rms`` it is the raw RMSNorm input. + b: Weight [K, N], column-major (``W.t()`` of a contiguous [N, K]); with ``row_rms`` + the RMSNorm weight is already folded in. + out: [M, N] contiguous: BF16, or uint8 (E4M3 bits) with ``epilogue.fp8_rope``. + pdl: Launch with Programmatic Dependent Launch. + tactic: The split-K tactic; its A/B stages are reduced to fit the epilogue buffers. + epilogue: The fused epilogue. + gamma, norm_out: ``rms_out`` operands (BF16 [K], BF16 [M, K]). + kv_lens, cos_sin, q_scale: ``fp8_rope`` operands (int32 KV lengths per sequence, the + FP32 (cos, sin) table of ``mla_rope_generation``, FP32 [1] scale). + reserve_full_smem: Reserve the full shared-memory capacity per CTA. + """ + m, n, k = _validate_fused(a, b, out, epilogue, gamma, norm_out, kv_lens, cos_sin, q_scale) + tactic = fused_epilogue_tactic(tactic) + validate_tactic(tactic, m, n, k) + _check( + _fused_smem_bytes(tactic, tactic.ab_stages) <= _SMEM_CAPACITY_BYTES, + f"tactic {tactic} does not fit the epilogue buffers", + ) + tensors, leading_dims = _swap_cute_tensors(a, b, out) + device_index = a.device.index if a.device.index is not None else _torch.cuda.current_device() + compiled = _get_compiled_fused_kernel( + tactic, pdl, epilogue, leading_dims, out.dtype, reserve_full_smem, device_index + ) + ph = _placeholders(device_index) + aux = _aux_cute_tensors( + gamma if epilogue.rms_out else ph["gamma"], + norm_out if epilogue.rms_out else ph["norm_out"], + kv_lens if epilogue.fp8_rope else ph["kv_lens"], + cos_sin.view(-1) if epilogue.fp8_rope else ph["cos_sin"], + q_scale if epilogue.fp8_rope else ph["q_scale"], + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + compiled(*tensors, *aux, stream) + return out + + +@cute.experimental.jit +def _bmm_qtail( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + positions: cute.Tensor, + cos_sin: cute.Tensor, + weights_in: cute.Tensor, + q_scale_out: cute.Tensor, + weights_out: cute.Tensor, + stream: _cuda.CUstream, +): + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])), + positions, + cos_sin, + weights_in, + q_scale_out, + weights_out, + stream, + ) + + +def _qtail_cute_tensors(positions, cos_sin, weights_in, q_scale_out, weights_out): + return ( + _from_dlpack_dynamic(positions.detach(), 0, assumed_align=4), + _from_dlpack_dynamic(cos_sin.detach(), 1, assumed_align=4), + _from_dlpack_dynamic(weights_in.detach(), 1, assumed_align=2), + _from_dlpack_dynamic(q_scale_out.detach(), 1, assumed_align=4), + _from_dlpack_dynamic(weights_out.detach(), 1, assumed_align=4), + ) + + +@functools.cache +def _get_compiled_qtail_kernel( + tactic: SplitKTactic, + use_pdl: bool, + epilogue: SplitKEpilogue, + qtail: IndexerQTail, + leading_dims: tuple[int, int, int], + reserve_full_smem: bool, + device_index: int, +): + kernel = SplitKIndexerQTailGemmKernel( + tactic=tactic, + use_pdl=use_pdl, + epilogue=epilogue, + qtail=qtail, + reserve_full_smem=reserve_full_smem, + ) + device = _torch.device("cuda", device_index) + m, n, k = 8, qtail.n_heads * qtail.head_dim, epilogue.rms_k + weight = _torch.empty((n, k), dtype=_torch.bfloat16, device=device) + act = _torch.empty((m, k), dtype=_torch.bfloat16, device=device) + out = _torch.empty((m, n), dtype=_torch.uint8, device=device) + tensors, repr_leading = _swap_cute_tensors(act, weight.t(), out) + if repr_leading != leading_dims: + raise ValueError(f"unsupported operand layouts {leading_dims} (expected {repr_leading})") + aux = _qtail_cute_tensors( + _torch.zeros(m, dtype=_torch.int32, device=device), + _torch.zeros((m, qtail.rope_dim), dtype=_torch.float32, device=device), + _torch.zeros((m, qtail.n_heads), dtype=_torch.bfloat16, device=device), + _torch.zeros((m, qtail.n_heads), dtype=_torch.float32, device=device), + _torch.zeros((m, qtail.n_heads), dtype=_torch.float32, device=device), + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(device).cuda_stream) + return cute_ext.compile(_bmm_qtail, kernel, *tensors, *aux, stream) + + +def _validate_qtail( + a, b, out, epilogue, qtail, positions, cos_sin, weights_in, q_scale_out, weights_out +): + _check(all(isinstance(t, _torch.Tensor) for t in (a, b, out)), "a, b and out must be tensors") + _check(a.ndim == 2 and b.ndim == 2 and out.ndim == 2, "a, b and out must be 2D") + _check( + a.is_cuda and b.device == a.device and out.device == a.device, + "a, b and out must be on one CUDA device", + ) + _check(a.dtype == _torch.bfloat16 and b.dtype == _torch.bfloat16, "a and b must be BF16") + m, k = a.shape + _check(b.shape[0] == k, f"incompatible shapes {tuple(a.shape)} x {tuple(b.shape)}") + n = b.shape[1] + _check(tuple(out.shape) == (m, n), f"out must be {(m, n)}, got {tuple(out.shape)}") + _check(1 <= m <= 32 and k % _CTA_K == 0, f"unsupported shape M={m} N={n} K={k}") + _check( + a.stride(1) == 1 and a.stride(0) >= k and (a.stride(0) * 2) % 16 == 0, + f"a needs unit inner stride and a 16-byte row pitch, got strides {a.stride()}", + ) + _check(b.t().is_contiguous(), "b must be column-major (the transpose of a contiguous [N, K])") + _check(out.is_contiguous() and out.dtype == _torch.uint8, "out must be contiguous uint8") + _check(all(t.data_ptr() % 32 == 0 for t in (a, b, out)), "a, b and out must be 32-byte aligned") + _check( + epilogue.row_rms and not epilogue.rms_out and not epilogue.fp8_rope, + "the indexer Q tail follows the row_rms epilogue alone", + ) + _check(k == epilogue.rms_k, f"row_rms was specialized for K={epilogue.rms_k}, got {k}") + qtail.validate() + _check( + n == qtail.n_heads * qtail.head_dim, + f"N={n} is not {qtail.n_heads} heads of {qtail.head_dim}", + ) + _check( + isinstance(positions, _torch.Tensor) + and positions.dtype == _torch.int32 + and positions.dim() == 1 + and positions.is_contiguous() + and positions.numel() >= m, + "positions must be a contiguous int32 vector of at least M entries", + ) + _check( + isinstance(cos_sin, _torch.Tensor) + and cos_sin.dtype == _torch.float32 + and cos_sin.dim() == 2 + and cos_sin.stride(1) == 1 + and cos_sin.shape[1] >= qtail.rope_dim, + "cos_sin must be FP32 [positions, >= rope_dim] with unit inner stride", + ) + _check( + isinstance(weights_in, _torch.Tensor) + and weights_in.dtype == _torch.bfloat16 + and tuple(weights_in.shape) == (m, qtail.n_heads) + and weights_in.stride(1) == 1, + "weights_in must be BF16 [M, n_heads] with unit inner stride", + ) + for t, name in ((q_scale_out, "q_scale_out"), (weights_out, "weights_out")): + _check( + isinstance(t, _torch.Tensor) + and t.dtype == _torch.float32 + and tuple(t.shape) == (m, qtail.n_heads) + and t.is_contiguous(), + f"{name} must be a contiguous FP32 [M, n_heads]", + ) + for t in (positions, cos_sin, weights_in, q_scale_out, weights_out): + _check(t.device == a.device, "every operand must be on the GEMM's device") + return m, n, k + + +def run_splitk_idx_qtail( + a: _torch.Tensor, + b: _torch.Tensor, + out: _torch.Tensor, + pdl: bool, + tactic: SplitKTactic, + epilogue: SplitKEpilogue, + qtail: IndexerQTail, + *, + positions: _torch.Tensor, + cos_sin: _torch.Tensor, + weights_in: _torch.Tensor, + q_scale_out: _torch.Tensor, + weights_out: _torch.Tensor, + reserve_full_smem: bool = True, +) -> _torch.Tensor: + """The DSA indexer wq_b ``a[M, K] @ b[K, N]`` with the row-RMS epilogue and the indexer Q tail. + + Args: + a: The raw q_a [M, K] BF16 with unit inner stride (may be a row view of a wider buffer). + b: The norm-folded indexer wq_b [K, N], column-major (``W.t()`` of a contiguous [N, K]). + out: [M, N] uint8: the E4M3 indexer Q, [tokens, heads * head_dim]. + pdl: Launch with Programmatic Dependent Launch. + tactic: The split-K tactic: ``mma_m`` is one head or half a head; its A/B stages are + reduced to fit the epilogue buffers. + epilogue: ``row_rms`` alone. + qtail: The Q tail. + positions: int32 [>= M] RoPE positions of the tokens. + cos_sin: FP32 [positions, >= rope_dim] cos/sin rows ([cos | sin] halves). + weights_in: BF16 [M, n_heads] attention weights (the weights_proj output; may be a column + view of a wider buffer). + q_scale_out: FP32 [M, n_heads] output: the E4M3 scales. + weights_out: FP32 [M, n_heads] output: weights_in * scale * qtail.weight_scale. + reserve_full_smem: Reserve the full shared-memory capacity per CTA. + """ + m, n, k = _validate_qtail( + a, b, out, epilogue, qtail, positions, cos_sin, weights_in, q_scale_out, weights_out + ) + tactic = idx_qtail_tactic(tactic, qtail) + validate_tactic(tactic, m, n, k) + _check( + _qtail_smem_bytes(tactic, tactic.ab_stages, qtail) <= _SMEM_CAPACITY_BYTES, + f"tactic {tactic} does not fit the indexer Q tail buffers", + ) + tensors, leading_dims = _swap_cute_tensors(a, b, out) + device_index = a.device.index if a.device.index is not None else _torch.cuda.current_device() + compiled = _get_compiled_qtail_kernel( + tactic, pdl, epilogue, qtail, leading_dims, reserve_full_smem, device_index + ) + aux = _qtail_cute_tensors(positions, cos_sin, weights_in, q_scale_out, weights_out) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + compiled(*tensors, *aux, stream) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/dense_mlp_small_m.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/dense_mlp_small_m.py new file mode 100644 index 000000000000..577c93008ffd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/dense_mlp_small_m.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One rank's slice of the BF16 dense MLP for up to 8 decode rows, via the trtllm +glm_dense_mlp_small_m op (two weight-streaming kernels).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +#: GLM-5.2's hidden size, the contraction of gate/up and the output width of down. +HIDDEN = 6144 +#: The rank-local intermediate size of the dense layers 0-2 at TP16: 12288 / 16. +INTER_PER_RANK = 768 + + +class GlmDenseMlpSmallM(OpWrapper): + """`out = (silu(x @ W_gate.T) * (x @ W_up.T)) @ W_down.T` on one rank's slice, 1-8 rows. + + The slice is TP16's: `gate_up_weight` [2 * 768, 6144] stacks this rank's + 768 gate_proj rows over the same 768 up_proj rows (gate first), and + `down_weight` [6144, 768] holds the same 768 input columns of down_proj. + `out` [M, 6144] is the rank-local partial: the sum over the 16 ranks, which + the post-MLP all-reduce forms, is the MLP output. Returns a new tensor; no + input is modified. + + Rounding points are the unfused chain's: the gate and up sums round to + bf16 (the bf16 gate_up GEMM output), SwiGLU runs in fp32 and rounds once to + bf16 `h`, and the down sum rounds to bf16. Sums are fp32 in another order. + + Two kernels on the current stream: gate/up + SwiGLU (192 CTAs in 2-CTA + clusters that split the 6144 contraction) and down (96 CTAs of 64 output + columns). Under programmatic dependent launch each CTA bulk-copies all of + its weights into shared memory before its grid-dependency wait, and reads + `x` (gate/up) or `h` (down) only after it. gate/up releases down at entry; + down releases its own dependents once its weight copies are issued, so the + kernel after the op must wait before it reads `out`. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The target calls the op on the dense layers 0-2 of a generation-only + # forward: rows = requests x (draft length + 1), at most 8. Operand scales + # are the checkpoint's: post_attention_layernorm's gamma averages 0.021 on + # these layers (x has unit RMS before it), gate/up weights 0.014 std, down + # weights 0.013 std. `saturated` scales gate/up so the gate sums reach the + # sigmoid's tails. + CELLS: tuple[Cell, ...] = ( + Cell( + why="c1 verify forward: MTP K=7's 8 rows through a TP16 dense layer", + spec=dict(rows=8, x_scale=0.021, gate_up_std=0.014, down_std=0.013), + ), + Cell( + why="decode without speculation: 1 row, 7 of the 8 MMA token columns empty", + spec=dict(rows=1, x_scale=0.021, gate_up_std=0.014, down_std=0.013), + ), + Cell( + why="MTP K=3: 4 rows, the empty token columns start at a thread's pair", + spec=dict(rows=4, x_scale=0.021, gate_up_std=0.014, down_std=0.013), + ), + Cell( + why="saturated SwiGLU: gate sums to +-40, the fast exp and division in sigmoid's tails", + spec=dict(rows=8, x_scale=1.0, gate_up_std=0.13, down_std=0.036), + ), + ) + + note = """ + Shapes: hidden 6144, rank-local intermediate 768, 1-8 rows; BF16, + contiguous, 16-byte aligned operands on an sm_90 or newer device. The op + raises on anything else, so the caller routes other shapes (more rows, + context forwards, another TP) to a generic gated MLP. + `torch.ops.trtllm.glm_dense_mlp_small_m_supported(rows, hidden, inter)` + answers the shape question on the host. + + The weight layout is the caller's to get right and no check can see it: a + gate_up_weight that stacks up over gate, or that takes rows + [r * 1536, (r + 1) * 1536) of a pre-fused [gate; up] tensor, gives a + plausible wrong MLP. + + The weights are read before the grid-dependency wait, so they must not be + written by any kernel still running ahead of the op on the stream. + """ + + def __call__( + self, x: torch.Tensor, gate_up_weight: torch.Tensor, down_weight: torch.Tensor + ) -> torch.Tensor: + return torch.ops.trtllm.glm_dense_mlp_small_m(x, gate_up_weight, down_weight) + + def reference( + self, x: torch.Tensor, gate_up_weight: torch.Tensor, down_weight: torch.Tensor + ) -> torch.Tensor: + """The rank partial with the op's rounding points, fp32 sums. + + The intermediate bf16 roundings are part of what the op computes (the + unfused bf16 GEMMs round there too), so the reference rounds at the + same points instead of carrying fp32 through; without them the band + would measure those roundings, not the kernel. + + TF32 off: torch's default fp32 matmul on this hardware is a TF32 + product, which would make the reference the inaccurate side. + """ + inter = down_weight.shape[1] + with true_fp32_matmul(): + gate_up = x.float() @ gate_up_weight.float().t() + gate = gate_up[:, :inter].to(torch.bfloat16).float() + up = gate_up[:, inter:].to(torch.bfloat16).float() + h = (torch.nn.functional.silu(gate) * up).to(torch.bfloat16) + with true_fp32_matmul(): + return (h.float() @ down_weight.float().t()).to(x.dtype) + + def is_valid( + self, x: torch.Tensor, gate_up_weight: torch.Tensor, down_weight: torch.Tensor + ) -> None: + """Nothing to refuse: every input the op accepts it computes as stated. + + The op raises on every shape, dtype, layout, alignment and device it + does not serve, including pre-sm_90 devices where the kernels have no + body. What is left is the weight layout (see `note`), which no check on + the tensors can see. + """ + return None + + +dense_mlp_small_m = GlmDenseMlpSmallM() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/draft_hidden_half.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/draft_hidden_half.py new file mode 100644 index 000000000000..99bdabd41f32 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/draft_hidden_half.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""This rank's FP32 share of the next MTP draft step's eh_proj hidden half, hnorm included, as +one split-K GEMM.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +from ._low_m_splitk import ( + device_limits, + operands_overlap, + run_splitk_hidden_half, + splitk_default_tactic, +) + +#: GLM-5.2 at tp16: hidden 6144, this rank's 6144 / 16 = 384 output rows of eh_proj's hidden +#: half, hnorm eps 1e-5. +_GLM = dict(k=6144, n=384, eps=1e-5) + + +class GlmDraftHiddenHalf(OpWrapper): + """`out = rmsnorm(x) @ weight.T` (with `norm_weight`) or `out = x @ weight.T` (without), FP32 + out; returns `out`. + + eh_proj is row-parallel over `cat(enorm(embed(token)), hnorm(hidden))`, so its output is a + token half (tabulated at load) plus `hnorm(hidden) @ W_e[:, H:].T`. This op computes this + rank's `H / tp` outputs of that hidden half in FP32, on the draft step's hidden state before + the draft LM head runs. `weight` is this rank's rows `[rank * H / tp, (rank + 1) * H / tp)` + of `W_e[:, H:]` (gathered from the row-parallel shards once, at load). + + With `norm_weight` (one row of H = 6144): the kernel's B producer computes FlashInfer's + RMSNorm of `x` itself -- the 64-thread layout, the sequential square sums, the two + butterflies and their sum, `rsqrt.approx.ftz` of the IEEE `sum / H + eps`, `(x * rstd) * + (w + 0)` rounded to BF16 -- so the result is bit-identical to running that RMSNorm first and + the GEMM on its output. Without it `x` is the normalized row(s). + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="MTP draft steps 0-5 (one request): 1 row, hnorm in the kernel, split-K 8 " + "(48 CTAs, every K tile in its own stage)", + spec=dict(_GLM, rows=1, in_kernel_norm=True), + ), + Cell( + why="two requests per draft step: 2 normalized rows, the plain FP32-output GEMM", + spec=dict(_GLM, rows=2, in_kernel_norm=False), + ), + ) + + note = """ + The tactic is the shape's `splitk_default_tactic`; the in-kernel norm needs one row, H = + 6144 and a tactic whose CTA K slice fits its stages (`hnorm_b_supported`), else ValueError. + PDL follows `TRTLLM_ENABLE_PDL` (default on): before its grid dependency wait the kernel + reads only the weight (TMA) and the norm weight and zeroes its own B stages; it reads `x` + after the wait. + + Loud, so not guarded (ValueError before the launch): M outside [1, 32], K not a multiple of + 128, non-BF16 x or weight, a non-FP32 or non-contiguous out, pointers not 32-byte aligned, + the in-kernel norm where it does not serve. The first call of a (tactic, norm eps) compiles + the kernel; make it before CUDA-graph capture. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + norm_weight: torch.Tensor | None, + eps: float | None, + ) -> torch.Tensor: + """Write `out` and return it. + + Args: + x: [rows, H] BF16 with unit inner stride: the raw hidden state with `norm_weight`, + else the normalized one. + weight: [N, H] BF16, contiguous: this rank's hidden-half rows of eh_proj. + out: [rows, N] FP32, contiguous. + norm_weight: [H] BF16 hnorm weight for the in-kernel norm, or None. + eps: hnorm's epsilon (with `norm_weight`), else None. + """ + rows, k = x.shape + sm_count, smem_capacity = device_limits(x.device.index) + tactic = splitk_default_tactic( + rows, weight.shape[0], k, sm_count=sm_count, smem_capacity=smem_capacity + ) + return run_splitk_hidden_half( + x, weight, out, get_env_enable_pdl(), tactic, norm_weight=norm_weight, norm_eps=eps + ) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + norm_weight: torch.Tensor | None, + eps: float | None, + ) -> torch.Tensor: + """FP32 product (TF32 off) of the BF16 normalized row(s) and the weight. The norm is + FlashInfer's formula in FP32 -- `rstd = rsqrt(sum(x ** 2) / H + eps)`, `(x * rstd) * w` + rounded to BF16 once -- with torch's summation order. `out` is not read.""" + xf = x.float() + if norm_weight is not None: + rstd = torch.rsqrt(xf.pow(2).sum(dim=-1, keepdim=True) / xf.shape[-1] + eps) + xf = ((xf * rstd) * norm_weight.float()).to(torch.bfloat16).float() + with true_fp32_matmul(): + return xf @ weight.float().t() + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + norm_weight: torch.Tensor | None, + eps: float | None, + ) -> None: + # `eps` without `norm_weight` is ignored: the row goes into the GEMM unnormalized. + assert (norm_weight is None) == (eps is None), ( + "pass norm_weight and eps together: eps alone leaves x unnormalized" + ) + # The owner CTA of a tile stores while other clusters still read x and the weight. + for operand, name in ((x, "x"), (weight, "weight")): + assert not operands_overlap(out, operand), f"out must not share memory with {name}" + + +draft_hidden_half = GlmDraftHiddenHalf() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/fused_a_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/fused_a_gemm.py new file mode 100644 index 000000000000..5f02dd788c34 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/fused_a_gemm.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GLM-5 fused-A GEMM (kv_a_proj_with_mqa, optionally with the DSA indexer's wk | weights_proj +rows) for up to 8 tokens, with an L2 prefetch of later weights, via the trtllm glm_fused_a_gemm +op.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +HIDDEN_SIZE = 6144 +# kv_a_proj_with_mqa: q_a (2048) | compressed KV (512) | k_pe (64). +OUT_FEATURES = 2624 +# The same rows, then the indexer's wk (128) and weights_proj (32) rows. +INDEXER_OUT_FEATURES = 2784 +ROWS_PER_CTA = 16 +MAX_TOKENS = 8 + +# What a c1 decode step lists behind fused-A, at TP16 per-rank shapes. +_INDEXER_WQ_B = ((4096, 2048), torch.bfloat16) # indexer wq_b, q_a norm folded in (not TP-split) +_Q_B_ABSORBED = ((2304, 2048), torch.bfloat16) # q_b with k_b absorbed: 4 heads x (512 + 64) +_V_B = ((4, 256, 512), torch.bfloat16) # v_b_proj, 4 heads +_O_PROJ = ((6144, 1024), torch.bfloat16) # o_proj shard +_LEADER_TOPK = ((8, 2048), torch.int32) # the indexer layer's top-k that a shared layer reuses + + +class GlmFusedAGemm(OpWrapper): + """`hidden_states @ weight.T`: BF16 operands, FP32 accumulation, BF16 out, `[tokens, rows]`. + + `weight` is the kv_a_proj_with_mqa weight as stored, `[2624, 6144]` row-major (q_a, compressed + KV, k_pe), or one tensor with the DSA indexer's wk and weights_proj rows appended, + `[2784, 6144]`. Output columns then read `[0, 2048)` q_a, `[2048, 2560)` compressed KV, + `[2560, 2624)` k_pe, `[2624, 2752)` indexer K, `[2752, 2784)` indexer weights; columns + `[0, 2624)` are bitwise those of the 2624-row call. + + One CTA per 16 weight rows (164 or 174 CTAs, one per SM) with all eight warps issuing the + HMMAs. Before its grid-dependency wait the kernel stages the CTA's weight rows (tensor copies + into shared memory, two thirds of them on into registers) and then prefetches the + `l2_prefetch` tensors into L2; it reads `hidden_states` only after the wait and triggers its + dependents right after the wait. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The c1 decode step runs the op 85 times: 78 verify layers at 8 rows plus 7 MTP draft steps. + # 22 calls (21 indexer layers and draft step 0) take the 2784-row weight. + CELLS: tuple[Cell, ...] = ( + Cell( + why="indexer layers' verify step and MTP draft step 0: 2784 rows, 8 tokens (3D " + "activation copies), the list indexer wq_b, v_b, o_proj", + spec=dict(tokens=8, out_features=2784, l2_prefetch=(_INDEXER_WQ_B, _V_B, _O_PROJ)), + ), + Cell( + why="shared-top-k layers' verify step (57 of 78): 2624 rows, 8 tokens, a full list " + "of four ending in the indexer layer's int32 top-k", + spec=dict( + tokens=8, + out_features=2624, + l2_prefetch=(_Q_B_ABSORBED, _V_B, _O_PROJ, _LEADER_TOPK), + ), + ), + Cell( + why="MTP draft steps 1-6: 2624 rows at one token (2D activation boxes of one row)", + spec=dict(tokens=1, out_features=2624, l2_prefetch=(_Q_B_ABSORBED, _V_B, _O_PROJ)), + ), + Cell( + why="4 tokens (MTP with 3 draft tokens): 2D boxes between 1 and 8 rows, 2784 rows", + spec=dict(tokens=4, out_features=2784, l2_prefetch=(_INDEXER_WQ_B, _V_B, _O_PROJ)), + ), + Cell( + why="empty list: the instance without L2 prefetches, for more than 32 rows per step", + spec=dict(tokens=8, out_features=2624, l2_prefetch=()), + ), + ) + + note = """ + `l2_prefetch` is a cache hint: up to four contiguous, 16-byte aligned CUDA tensors on the + GEMM's device, issued in 16 KB bulk prefetches spread over the CTAs once each CTA's weight + rows are staged. It never enters the output, and the kernel writes none of it. + + PDL (`TRTLLM_ENABLE_PDL`): before its grid-dependency wait the kernel reads `weight`, so no + kernel launched earlier in the step may write it (load-time weights only). Its dependents + launch once it passes the wait and must wait on their own grid dependency before reading the + output or writing `hidden_states`. + + Preconditions, enforced loudly by the op: CUDA BF16 2-D `hidden_states` and `weight` on one + device, contiguous and 16-byte aligned; hidden 6144; 2624 or 2784 rows; 1 to 8 tokens; at most + four `l2_prefetch` tensors, each contiguous, aligned and on the device; SM90 or later. + + Numerics: the products of a plain GEMM, summed in FP32 in the kernel's order (two + accumulators per warp, the eight warp partials in warp order) and rounded to BF16 once; + deterministic. The shared `dsv3_fused_a_gemm_op` sums the same products in another order: + nearly all outputs agree bitwise, the rest by one BF16 rounding step. + """ + + def __call__( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + l2_prefetch: list[torch.Tensor], + ) -> torch.Tensor: + return torch.ops.trtllm.glm_fused_a_gemm(hidden_states, weight, l2_prefetch) + + def reference( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + l2_prefetch: list[torch.Tensor], + ) -> torch.Tensor: + """`hidden_states @ weight.T` in FP32 (TF32 off), rounded to BF16 once. + + `l2_prefetch` does not enter the result; it keeps its name to mirror `__call__`. + """ + with true_fp32_matmul(): + return (hidden_states.float() @ weight.float().t()).to(hidden_states.dtype) + + def is_valid( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + l2_prefetch: list[torch.Tensor], + ) -> None: + # The shape gates of the kernel, which a target checks to route anything else to the + # shared dsv3_fused_a_gemm_op. The op raises on all but the last of them as well. + # gemm_k 6144: 24 k stages of 256, three per MMA warp. + assert hidden_states.dim() == 2 and hidden_states.shape[1] == HIDDEN_SIZE, ( + f"hidden_states must be [tokens, {HIDDEN_SIZE}]" + ) + # Tiles of 16 weight rows, one CTA each; instances exist for the 2624-row weight and the + # 2784-row weight with the indexer rows. + assert weight.dim() == 2 and weight.shape[1] == HIDDEN_SIZE, ( + f"weight must be [rows, {HIDDEN_SIZE}]" + ) + assert weight.shape[0] in (OUT_FEATURES, INDEXER_OUT_FEATURES), ( + f"weight must have {OUT_FEATURES} or {INDEXER_OUT_FEATURES} rows, got {weight.shape[0]}" + ) + # One activation tile of 8 token rows per k stage. + assert 1 <= hidden_states.shape[0] <= MAX_TOKENS, ( + f"the kernel takes 1 to {MAX_TOKENS} tokens, got {hidden_states.shape[0]}" + ) + # One CTA per SM in a single wave: the kernel stages its weights before the + # grid-dependency wait, and a second wave of CTAs would start only after the first ends. + # The output stays right; on fewer SMs the caller takes the shared dsv3_fused_a_gemm_op + # instead (32-row tiles, half the CTAs, for the 2624-row weight). + sms = torch.cuda.get_device_properties(hidden_states.device).multi_processor_count + ctas = weight.shape[0] // ROWS_PER_CTA + assert sms >= ctas, f"{ctas} CTAs need as many SMs for one wave, the device has {sms} SMs" + + +fused_a_gemm = GlmFusedAGemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/indexer_q_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/indexer_q_tail.py new file mode 100644 index 000000000000..a8c071f53b76 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/indexer_q_tail.py @@ -0,0 +1,248 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DSA indexer wq_b from the raw q_a as one split-K GEMM whose epilogue finishes the indexer Q: +RoPE, the per-head E4M3 scale, the E4M3 Q and the scaled head weights.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +from ._low_m_splitk import operands_overlap +from ._low_m_splitk_fused import IndexerQTail, SplitKEpilogue, idx_qtail_plan, run_splitk_idx_qtail + +#: GLM-5.2's DSA indexer (replicated on every rank): 32 heads of 128, RoPE on the first 64 +#: columns as GPT-J pairs, UE8M0 scales, weight scale 128 ** -0.5 * 32 ** -0.5; q_lora_rank 2048 +#: and the q_a RMSNorm eps 1e-5. +_GLM = dict( + n_heads=32, + head_dim=128, + rope_dim=64, + interleave=True, + ue8m0=True, + weight_scale=128**-0.5 * 32**-0.5, + k=2048, + eps=1e-5, +) + +#: The E4M3 ratio of the indexer prologue: max(amax, 1e-12) * fp32(1 / 448). +_AMAX_FLOOR = 1e-12 +_INV_E4M3_MAX = 1.0 / 448.0 + + +class GlmIndexerQTail(OpWrapper): + """The DSA indexer's FP8 Q, Q scales and scaled head weights from the raw q_a; returns + `(q_fp8, q_scale, weights_out)`. + + The GEMM is the indexer wq_b with the q_a RMSNorm folded into its weight (`weight` is + `W_wq_b * gamma_q_a`, per column, built in FP64 and rounded to BF16 once at load): `y = bf16( + rstd_t * (q_a @ weight.T))`, `rstd_t = rsqrt(mean(q_a[t] ** 2) + eps)`, rounded to BF16 as a + BF16 store would. Then, per (token t, head h) of `y`, the indexer prologue's tail in its + order: RoPE on the head's first `rope_dim` columns (GPT-J pairs with `interleave`, NeoX + halves otherwise; cos/sin from the row of `positions[t]`) as `bf16(fma(x, cos, rn(+-partner * + sin)))`; `scale = max(amax, 1e-12) * fp32(1 / 448)`, rounded up to a power of two with + `ue8m0`; `q_fp8 = e4m3(q * (1 / scale))`; `q_scale = scale`; `weights_out = (weights_in * + scale) * weight_scale`. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="full-indexer layer, verify step and MTP draft step 0: 8 tokens, 32 heads of 128; " + "half-head tiles in (2, 2) clusters (128 CTAs)", + spec=dict(_GLM, m=8), + ), + Cell( + why="1 token, decode without MTP: one row in an 8-row tile (the same tactic)", + spec=dict(_GLM, m=1), + ), + ) + + note = """ + Kernel: `SplitKIndexerQTailGemmKernel`, the row-RMS fused split-K GEMM at the shape's + `splitk_default_tactic` (64-row tiles, split-K 2 for 32 x 128 outputs from 2048) whose + 64-row tiles hold half a head; the two halves of a head run in one cluster whose owner CTAs + exchange their row amax through DSMEM. The owner's epilogue warps compute the row RMS after + the grid dependency wait, then read the positions, the cos/sin rows and weights_in while the + mainloop runs. PDL follows `TRTLLM_ENABLE_PDL` (default on): before its wait the kernel reads + only the weight. + + The values are those of the fused indexer prologue run on wq_b's BF16 output: the same + operations in the same order. The interleaved RoPE and UE8M0 scales are GLM-5.2's; the NeoX + form and plain FP32 scales are implemented but not cells. + + Loud, so not guarded (ValueError before the launch): M outside [1, 32], K not a multiple of + 128, N not n_heads * head_dim, a head_dim that is neither one nor two 64-row tiles, operands + of the wrong dtype, shape or stride, a rope_dim the epilogue cannot split. The first call of + a (tactic, Q tail) compiles the kernel; make it before CUDA-graph capture. + """ + + def __call__( + self, + q_a: torch.Tensor, + weight: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + weights_in: torch.Tensor, + q_fp8: torch.Tensor, + q_scale: torch.Tensor, + weights_out: torch.Tensor, + eps: float, + n_heads: int, + head_dim: int, + rope_dim: int, + interleave: bool, + ue8m0: bool, + weight_scale: float, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Write `q_fp8`, `q_scale` and `weights_out` and return them. + + Args: + q_a: [M, q_lora_rank] BF16, the raw q_a (before its RMSNorm); unit inner stride, a + row pitch of a multiple of 8 elements (a row view of the fused-A output works). + weight: [n_heads * head_dim, q_lora_rank] BF16, contiguous: wq_b, norm-folded. + positions: [>= M] int32 RoPE positions, contiguous. + cos_sin: [table positions, >= rope_dim] FP32 with unit inner stride: per position + the rope_dim / 2 cosines, then the rope_dim / 2 sines. + weights_in: [M, n_heads] BF16 head weights (weights_proj output) with unit inner + stride (a column view of the fused-A output works). + q_fp8: [M, n_heads, head_dim] float8_e4m3fn, contiguous (output). + q_scale: [M, n_heads, 1] FP32, contiguous (output). + weights_out: [M, n_heads] FP32, contiguous (output). + eps: The q_a RMSNorm epsilon. + n_heads, head_dim, rope_dim, interleave, ue8m0, weight_scale: The indexer's Q tail. + """ + m, k = q_a.shape + qtail = IndexerQTail( + n_heads=n_heads, + head_dim=head_dim, + rope_dim=rope_dim, + interleave=interleave, + ue8m0=ue8m0, + weight_scale=float(weight_scale), + ) + tactic = idx_qtail_plan(qtail, m, weight.shape[0], k, q_a.device.index) + run_splitk_idx_qtail( + q_a, + weight.t(), + q_fp8.view(torch.uint8).view(m, -1), + get_env_enable_pdl(), + tactic, + SplitKEpilogue(row_rms=True, rms_k=k, rms_eps=float(eps)), + qtail, + positions=positions, + cos_sin=cos_sin, + weights_in=weights_in, + q_scale_out=q_scale.view(m, n_heads), + weights_out=weights_out, + ) + return q_fp8, q_scale, weights_out + + def reference( + self, + q_a: torch.Tensor, + weight: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + weights_in: torch.Tensor, + q_fp8: torch.Tensor, + q_scale: torch.Tensor, + weights_out: torch.Tensor, + eps: float, + n_heads: int, + head_dim: int, + rope_dim: int, + interleave: bool, + ue8m0: bool, + weight_scale: float, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """`(q_fp8, q_scale, weights_out)` as new tensors. The GEMM in FP32 (TF32 off); every + rounding the prologue applies, applied here; each RoPE fma in FP64, rounded to FP32 + once. The output arguments are not read.""" + m = q_a.shape[0] + x = q_a.float() + rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) + with true_fp32_matmul(): + y = ((x @ weight.float().t()) * rstd).to(torch.bfloat16).float() + y = y.view(m, n_heads, head_dim) + half = rope_dim // 2 + rows = cos_sin.reshape(cos_sin.shape[0], -1)[positions[:m].long(), :rope_dim] + cos = rows[:, :half].unsqueeze(1) + sin = rows[:, half:rope_dim].unsqueeze(1) + pe = y[..., :rope_dim] + if interleave: + first, second = pe[..., 0::2], pe[..., 1::2] + else: + first, second = pe[..., :half], pe[..., half:] + + def fma(a, b, c): + return (a.double() * b.double() + c.double()).float() + + rot_first = fma(first, cos, -second * sin).to(torch.bfloat16).float() + rot_second = fma(second, cos, first * sin).to(torch.bfloat16).float() + if interleave: + pe = torch.stack((rot_first, rot_second), dim=-1).flatten(-2) + else: + pe = torch.cat((rot_first, rot_second), dim=-1) + q = torch.cat((pe, y[..., rope_dim:]), dim=-1) + ratio = q.abs().amax(dim=-1, keepdim=True).clamp_min(_AMAX_FLOOR) * torch.tensor( + _INV_E4M3_MAX, dtype=torch.float32 + ) + scale = _ue8m0_round_up(ratio) if ue8m0 else ratio + q_ref = (q * (1.0 / scale)).clamp(-448.0, 448.0).to(torch.float8_e4m3fn) + weights_ref = (weights_in.float() * scale.squeeze(-1)) * torch.tensor( + weight_scale, dtype=torch.float32 + ) + return q_ref, scale, weights_ref + + def is_valid( + self, + q_a: torch.Tensor, + weight: torch.Tensor, + positions: torch.Tensor, + cos_sin: torch.Tensor, + weights_in: torch.Tensor, + q_fp8: torch.Tensor, + q_scale: torch.Tensor, + weights_out: torch.Tensor, + eps: float, + n_heads: int, + head_dim: int, + rope_dim: int, + interleave: bool, + ue8m0: bool, + weight_scale: float, + ) -> None: + # A position outside the cos/sin table reads past it: garbage, not an error. + used = positions[: q_a.shape[0]].long() + assert bool((used >= 0).all()) and int(used.max()) < cos_sin.shape[0], ( + "every position must index a row of the cos/sin table" + ) + # The epilogue stores while other clusters still read q_a, the weight and weights_in. + outputs = {"q_fp8": q_fp8, "q_scale": q_scale, "weights_out": weights_out} + operands = {"q_a": q_a, "weight": weight, "weights_in": weights_in} + for output_name, output in outputs.items(): + for name, operand in operands.items(): + assert not operands_overlap(output, operand), ( + f"{output_name} must not share memory with {name}" + ) + + +def _ue8m0_round_up(ratio: torch.Tensor) -> torch.Tensor: + """The positive FP32 `ratio` with its mantissa cleared and, if it was nonzero, its exponent + raised by one.""" + bits = ratio.view(torch.int32) + exponent = bits & 0x7F800000 + exponent = exponent + ((bits & 0x007FFFFF) != 0).to(torch.int32) * 0x00800000 + return exponent.view(torch.float32) + + +indexer_q_tail = GlmIndexerQTail() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/low_m_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/low_m_gemm.py new file mode 100644 index 000000000000..a232199870d6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/low_m_gemm.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-M BF16 projection on the target's split-K CuTe DSL GEMM, into a caller-owned output.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +from ._low_m_splitk import device_limits, operands_overlap, run_splitk_dense, splitk_default_tactic + + +class GlmLowMGemm(OpWrapper): + """`out = x @ weight.T`, BF16 in and out, FP32 accumulation rounded once; returns `out`. + + One launch of the low-M split-K kernel at the shape's `splitk_default_tactic`: each CTA of a + cluster accumulates an exact K slice in FP32, the cluster reduces through DSMEM in a fixed + order, and the owner CTA rounds and stores. The decode step's plain BF16 projections of at + most 32 rows take it: the shared expert's gate_up and down, the LM head on this rank's + vocabulary shard, and MTP draft step 0's eh_proj on this rank's input slice. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # GLM-5.2 at tp16: hidden 6144, shared-expert intermediate 2048 / 16 = 128 (gate_up fused, + # 256 rows), vocabulary 154880 / 16 = 9680, eh_proj input 2 * 6144 / 16 = 768. A verify + # step runs 8 rows (MTP K=7); the MTP draft steps run the shared expert and the LM head on + # 1 row, draft step 0's eh_proj on 8. + CELLS: tuple[Cell, ...] = ( + Cell( + why="shared expert gate_up, verify step: 8 x 6144 -> 256, split-K 8 (32 CTAs)", + spec=dict(m=8, k=6144, n=256), + ), + Cell( + why="shared expert gate_up, MTP draft step: 1 row", + spec=dict(m=1, k=6144, n=256), + ), + Cell( + why="shared expert down, verify step: 8 x 128 -> 6144, one K tile, no split (96 CTAs)", + spec=dict(m=8, k=128, n=6144), + ), + Cell( + why="shared expert down, MTP draft step: 1 row", + spec=dict(m=1, k=128, n=6144), + ), + Cell( + why="LM head vocabulary shard, verify step: 8 x 6144 -> 9680, partial last tile", + spec=dict(m=8, k=6144, n=9680), + ), + Cell( + why="LM head vocabulary shard, MTP draft step: 1 row", + spec=dict(m=1, k=6144, n=9680), + ), + Cell( + why="eh_proj, MTP draft step 0: 8 x 768 (row-parallel slice) -> 6144, split-K 2", + spec=dict(m=8, k=768, n=6144), + ), + ) + + note = """ + The tactic is a function of (M, N, K) and the device's SM count and opt-in shared memory + (`splitk_default_tactic`); the launch reserves the full shared-memory capacity, one CTA per + SM. PDL follows `TRTLLM_ENABLE_PDL` (default on): the weight is streamed before the grid + dependency wait, `x` is read after it, and the kernel releases its dependents once its weight + loads are issued. + + Loud, so not guarded: M outside [1, 32], K not a multiple of 128, mismatched shapes, dtypes + or devices, operands that are neither row- nor column-major, pointers not 32-byte aligned + (ValueError before the launch). The first call of a (shape class, tactic) compiles the + kernel; make it before CUDA-graph capture. + """ + + def __call__(self, x: torch.Tensor, weight: torch.Tensor, out: torch.Tensor) -> torch.Tensor: + """x [M, K] BF16 row-major; weight [N, K] BF16 (the stored weight); out [M, N] BF16.""" + m, k = x.shape + n = weight.shape[0] + sm_count, smem_capacity = device_limits(x.device.index) + tactic = splitk_default_tactic(m, n, k, sm_count=sm_count, smem_capacity=smem_capacity) + return run_splitk_dense(x, weight.t(), None, out, get_env_enable_pdl(), tactic) + + def reference(self, x: torch.Tensor, weight: torch.Tensor, out: torch.Tensor) -> torch.Tensor: + """The product the op writes into `out`: FP32-accumulated (TF32 off), rounded once. + + `out` is named to mirror `__call__` and not read. + """ + with true_fp32_matmul(): + return (x.float() @ weight.float().t()).to(x.dtype) + + def is_valid(self, x: torch.Tensor, weight: torch.Tensor, out: torch.Tensor) -> None: + # The owner CTA of a tile stores its rows while other clusters still stream x and the + # weight; an `out` that shares memory with either is overwritten under their reads and + # the result is silently wrong. + for operand, name in ((x, "x"), (weight, "weight")): + assert not operands_overlap(out, operand), f"out must not share memory with {name}" + + +low_m_gemm = GlmLowMGemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/o_proj_publish.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/o_proj_publish.py new file mode 100644 index 000000000000..af36c46b03b8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/o_proj_publish.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Attention o_proj whose split-K epilogue stores the BF16 output straight into this rank's slot of +the post-attention MNNVL one-shot all-reduce's current Lamport buffer.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +from ._low_m_splitk import ( + LamportPublish, + device_limits, + operands_overlap, + publish_tactic, + run_splitk_publish, + splitk_default_tactic, +) + +#: GLM-5.2 o_proj at tp16: 4 heads * v_head_dim 256 = 1024 -> hidden 6144, 16 ranks. +_GLM = dict(k=1024, n=6144, world_size=16) + +#: Lamport buffers of an MNNVL one-shot workspace, and the int32 words of its buffer flags. +_LAMPORT_BUFFERS = 3 +_FLAG_WORDS = 4 + + +class GlmOProjPublish(OpWrapper): + """`x @ weight.T` rounded to BF16, stored at `[token][rank][column]` of the all-reduce + workspace's current Lamport buffer instead of into an output; returns None. + + The current buffer is `lamport` (bytes) + `flags[0] * flags[2]`: the flags' current index + times their bytes per buffer, read once per CTA after the grid dependency wait (the previous + all-reduce rotated them). Value (t, c) lands at element `(t * world_size + rank) * N + c` of + that buffer, in BF16 with -0.0 written as +0.0 (-0.0 is the buffer's idle value): this + rank's shard as the one-shot kernel would have stored it. The all-reduce of this output must + then take the published shard (it neither loads nor stores its own) and poll every slot. The + owner CTA of each 64 x 8 tile stores 16-byte vectors of 8 columns through `lamport`, which is + the workspace's multicast address in production, so the slot fills on every rank. + + `out` gives the GEMM its output layout and is not written. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="verify step: 8 tokens, 1024 -> 6144, split-K 2 (192 CTAs), rank 5 of 16", + spec=dict(_GLM, m=8, rank=5, current=1), + ), + Cell( + why="MTP draft steps 1-6: 1 token, last rank, buffer index 2", + spec=dict(_GLM, m=1, rank=15, current=2), + ), + Cell( + why="verify step at 4 ranks: the slot arithmetic of another group size, buffer 0", + spec=dict(_GLM, m=8, rank=2, current=0, world_size=4), + ), + ) + + note = """ + The tactic is the shape's `splitk_default_tactic` with its A/B stages lowered to fit the FP32 + staging tile (`publish_tactic`; the c1 o_proj keeps all 12). PDL: the kernel is always + expected to run with `TRTLLM_ENABLE_PDL` on (the default) in front of the all-reduce; its + epilogue warps wait on the grid dependency, release the dependents (the all-reduce) right + after the wait and only then read the flags. With PDL off the flags are read without a wait, + which is ordered by the plain stream dependency. + + The Lamport workspace, its multicast mapping and its flags belong to the all-reduce; this op + reads the flags and writes one slot per token, nothing else. Loud, so not guarded + (ValueError before the launch): M outside [1, 32], K not a multiple of 128, N not a multiple + of 8, non-BF16 operands, a `lamport` that is not a 1-D BF16 view or `flags` that is not a + 1-D int32 view. The first call of a (tactic, world size, rank) compiles the kernel; make it + before CUDA-graph capture. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + lamport: torch.Tensor, + flags: torch.Tensor, + world_size: int, + rank: int, + ) -> None: + """Publish `x @ weight.T` into the current Lamport buffer. + + Args: + x: [M, K] BF16, contiguous: the attention output of this rank's heads. + weight: [N, K] BF16, contiguous: o_proj's row-parallel shard. + out: [M, N] BF16, contiguous, 32-byte aligned: layout only, not written. + lamport: 1-D BF16 view of the workspace's Lamport buffers (multicast address). + flags: 1-D int32 view of the workspace's buffer flags, 16-byte aligned: [0] the + current buffer index, [2] the bytes per buffer. + world_size: Ranks of the all-reduce group. + rank: This rank in the group. + """ + m, k = x.shape + n = weight.shape[0] + sm_count, smem_capacity = device_limits(x.device.index) + tactic = publish_tactic( + splitk_default_tactic(m, n, k, sm_count=sm_count, smem_capacity=smem_capacity), + smem_capacity=smem_capacity, + ) + run_splitk_publish( + x, + weight.t(), + out, + lamport, + flags, + get_env_enable_pdl(), + tactic, + LamportPublish(world_size=world_size, rank=rank, trigger_after_wait=True), + ) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + lamport: torch.Tensor, + flags: torch.Tensor, + world_size: int, + rank: int, + ) -> torch.Tensor: + """The [M, N] BF16 values the op stores at slots `t * world_size + rank` of the current + buffer: FP32-accumulated (TF32 off), rounded once, -0.0 as +0.0. Where they land is the + test's to check; `out`, `lamport` and `flags` are not read.""" + with true_fp32_matmul(): + values = (x.float() @ weight.float().t()).to(torch.bfloat16) + return torch.where(values == 0, torch.zeros_like(values), values) + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + lamport: torch.Tensor, + flags: torch.Tensor, + world_size: int, + rank: int, + ) -> None: + # A rank outside the group lands in another token's slots. + assert 0 <= rank < world_size, f"rank {rank} is outside a group of {world_size}" + # Every store goes through the flags' current buffer; slots past its end land in the + # next buffer (or past the workspace) without complaint. + assert flags.numel() >= _FLAG_WORDS and flags.data_ptr() % 16 == 0, ( + "flags must hold the four 16-byte aligned words the kernel reads" + ) + current, _, bytes_per_buffer, _ = (int(v) for v in flags[:_FLAG_WORDS].tolist()) + m, n = x.shape[0], weight.shape[0] + assert 0 <= current < _LAMPORT_BUFFERS, f"current buffer index {current} is out of range" + assert m * world_size * n * lamport.element_size() <= bytes_per_buffer, ( + "the Lamport buffer is smaller than the [token][rank][column] slots it must hold" + ) + assert (current + 1) * bytes_per_buffer <= lamport.numel() * lamport.element_size(), ( + "the current Lamport buffer runs past the end of `lamport`" + ) + # The epilogue stores while other clusters still read x and the weight. + for operand, name in ((x, "x"), (weight, "weight")): + assert not operands_overlap(lamport, operand), ( + f"lamport must not share memory with {name}" + ) + + +o_proj_publish = GlmOProjPublish() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/q_b_fp8_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/q_b_fp8_rope.py new file mode 100644 index 000000000000..44637c3ea359 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/q_b_fp8_rope.py @@ -0,0 +1,213 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MLA q_b of a decode step as one split-K GEMM: k_b absorbed, the q_a RMSNorm folded, and the +RoPE'd E4M3 attention Q written by the epilogue.""" + +from __future__ import annotations + +from typing import Literal + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +from ._low_m_splitk import operands_overlap +from ._low_m_splitk_fused import SplitKEpilogue, q_b_plan, run_splitk_fused + +#: GLM-5.2 MLA at tp16: 4 heads per rank, kv_lora_rank 512 + qk_rope_head_dim 64 per head, +#: q_lora_rank 2048, RMSNorm eps 1e-5. +_GLM = dict(heads=4, head_dim=576, rope_dim=64, k=2048, eps=1e-5) + +#: E4M3's largest finite magnitude. +_E4M3_MAX = 448.0 + + +class GlmQBFp8Rope(OpWrapper): + """The generation attention Q of a GLM MLA layer from the raw q_a, as E4M3 bits; returns `out`. + + `out[t, h * head_dim + c] = e4m3(q_scale * rope_t(rstd_t * (q_a @ weight.T))[t, h, c])`, where + `rstd_t = rsqrt(mean(q_a[t] ** 2) + eps)` and `rope_t` rotates the last `rope_dim` columns of + every `head_dim`-wide head as GPT-J pairs (2p, 2p + 1) by the (cos, sin) pair p of the token's + position, `kv_lens[s] - seq_len + i` for token i of sequence s (`seq_len` tokens each). The + conversion rounds to nearest even and saturates at +-448. Everything is FP32 from the GEMM's + accumulator to the conversion; nothing is rounded to BF16 on the way. + + `weight` is q_b with k_b absorbed and the q_a norm weight folded in, built once at load (see + the target's weights): for head h, rows [h * head_dim, h * head_dim + kv_lora_rank) are + `k_b_trans[h] @ W_q_b[h, :qk_nope]` (`k_b_trans` the MLA's [heads, kv_lora_rank, qk_nope] + k_b) and the next `rope_dim` rows are `W_q_b[h, qk_nope:]`, every column k times + `gamma_q_a[k]`, computed in FP64 and rounded to BF16 once. The output is the + `[tokens, heads, kv_lora_rank + rope_dim]` FP8 Q the sparse-MLA generation kernel reads. + + `plan` picks the launch, never the result's definition: `regular` (the shape's + `splitk_default_tactic`), `beside_topk` (128-row tiles, no split-K, three stages, for the + copy queued beside the DSA indexer top-k), `early_gather` (split-K 4 where the regular tactic + is split-K 8, beside a flash decode that gathers before its wait). The plans differ in the + FP32 summation order only. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="full-indexer layer, verify step and MTP draft step 0: 8 tokens of one request, " + "beside the top-k (128-row tiles, no split-K, 18 CTAs)", + spec=dict(_GLM, m=8, seq_len=8, plan="beside_topk"), + ), + Cell( + why="shared-indexer layer, verify step: 8 tokens, early-gather split-K 4 (72 CTAs)", + spec=dict(_GLM, m=8, seq_len=8, plan="early_gather"), + ), + Cell( + why="MTP draft steps 1-6: 1 token, regular split-K 8 (144 CTAs)", + spec=dict(_GLM, m=1, seq_len=1, plan="regular"), + ), + Cell( + why="two requests of 8 tokens: per-sequence positions, 16-row tiles (regular plan)", + spec=dict(_GLM, m=16, seq_len=8, plan="regular"), + ), + ) + + note = """ + Kernel: the split-K GEMM's mainloop and DSMEM reduction with the owner CTA's epilogue staged + in shared memory. The owner's epilogue warps compute the row RMS from global memory while the + mainloop runs (after the grid dependency wait) and read kv_lens, the cos/sin rows and + q_scale after the MMA. PDL follows `TRTLLM_ENABLE_PDL` (default on): before its wait the + kernel reads only the weight. + + Loud, so not guarded (ValueError before the launch): M outside [1, 32], K not a multiple of + 128, N not a multiple of head_dim, M not a multiple of seq_len, too few kv_lens, an `out` + that is not contiguous uint8, a weight that is not a contiguous [N, K], a q_a without unit + inner stride and a 16-byte row pitch, a geometry the epilogue rejects. The first call of a + (tactic, epilogue) compiles the kernel; make it before CUDA-graph capture. + """ + + def __call__( + self, + q_a: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kv_lens: torch.Tensor, + cos_sin: torch.Tensor, + q_scale: torch.Tensor, + eps: float, + head_dim: int, + rope_dim: int, + seq_len: int, + plan: Literal["regular", "beside_topk", "early_gather"], + ) -> torch.Tensor: + """Write `out` and return it. + + Args: + q_a: [M, q_lora_rank] BF16, the raw q_a (before its RMSNorm); unit inner stride, a + row pitch of a multiple of 8 elements (a row view of the fused-A output works). + weight: [heads * head_dim, q_lora_rank] BF16, contiguous: absorbed and norm-folded. + out: [M, heads * head_dim] uint8, contiguous: E4M3 bits. + kv_lens: [sequences] int32 KV lengths including this step's tokens. + cos_sin: [positions, rope_dim, 2] FP32 contiguous (cos, sin) pairs; position p's + pair i rotates columns 2i and 2i + 1 of every head's RoPE slice (the first + rope_dim / 2 pairs of a row are read). + q_scale: [1] FP32, the factor applied before the E4M3 conversion. + eps: The q_a RMSNorm epsilon. + head_dim: kv_lora_rank + rope_dim. + rope_dim: The RoPE width at the end of every head. + seq_len: Tokens per sequence; M is sequences * seq_len. + plan: The launch plan. + """ + m, k = q_a.shape + tactic = q_b_plan(plan, m, weight.shape[0], k, q_a.device.index) + epilogue = SplitKEpilogue( + row_rms=True, + rms_k=k, + rms_eps=float(eps), + fp8_rope=True, + head_dim=head_dim, + rope_dim=rope_dim, + seq_len=seq_len, + ) + return run_splitk_fused( + q_a, + weight.t(), + out, + get_env_enable_pdl(), + tactic, + epilogue, + kv_lens=kv_lens, + cos_sin=cos_sin, + q_scale=q_scale, + ) + + def reference( + self, + q_a: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kv_lens: torch.Tensor, + cos_sin: torch.Tensor, + q_scale: torch.Tensor, + eps: float, + head_dim: int, + rope_dim: int, + seq_len: int, + plan: Literal["regular", "beside_topk", "early_gather"], + ) -> torch.Tensor: + """The E4M3 bits the op writes, as uint8 [M, N]: FP32 throughout (TF32 off), one + conversion at the end. `out` and `plan` are not read.""" + m = q_a.shape[0] + x = q_a.float() + rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) + with true_fp32_matmul(): + q = (x @ weight.float().t()) * rstd + q = q.view(m, -1, head_dim) + nope = head_dim - rope_dim + token = torch.arange(m, device=q.device) + seq = token // seq_len + position = kv_lens.long()[seq] - seq_len + (token - seq * seq_len) + pairs = cos_sin.view(-1, rope_dim, 2)[position, : rope_dim // 2] + cos = pairs[..., 0].unsqueeze(1) + sin = pairs[..., 1].unsqueeze(1) + pe = q[..., nope:].reshape(m, q.shape[1], rope_dim // 2, 2) + even, odd = pe[..., 0], pe[..., 1] + rotated = torch.stack((cos * even - sin * odd, cos * odd + sin * even), dim=-1) + q = torch.cat((q[..., :nope], rotated.flatten(-2)), dim=-1) * q_scale.float()[0] + q = q.clamp(-_E4M3_MAX, _E4M3_MAX).to(torch.float8_e4m3fn) + return q.view(torch.uint8).reshape(m, -1) + + def is_valid( + self, + q_a: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kv_lens: torch.Tensor, + cos_sin: torch.Tensor, + q_scale: torch.Tensor, + eps: float, + head_dim: int, + rope_dim: int, + seq_len: int, + plan: Literal["regular", "beside_topk", "early_gather"], + ) -> None: + # The kernel indexes the flat table as (position * rope_dim + pair) * 2: a table of + # another row width is read at the wrong rows without complaint. + assert cos_sin.dim() == 3 and tuple(cos_sin.shape[1:]) == (rope_dim, 2), ( + f"cos_sin must be [positions, rope_dim={rope_dim}, 2], got {tuple(cos_sin.shape)}" + ) + # A position outside the table reads past it (or before it): garbage, not an error. + sequences = q_a.shape[0] // seq_len + lens = kv_lens[:sequences].long() + assert bool((lens >= seq_len).all()), "kv_lens must cover the step's own tokens" + assert int(lens.max()) <= cos_sin.shape[0], ( + "the cos/sin table must hold every position up to max(kv_lens) - 1" + ) + # The epilogue stores while other clusters still read q_a and the weight. + for operand, name in ((q_a, "q_a"), (weight, "weight")): + assert not operands_overlap(out, operand), f"out must not share memory with {name}" + + +q_b_fp8_rope = GlmQBFp8Rope() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/router_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/router_gemm.py new file mode 100644 index 000000000000..990fe197a7ea --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/router_gemm.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GLM-5 MoE router logits (256 experts at hidden 6144, FP32 from BF16) for up to 16 tokens, via +the trtllm glm_router_gemm op.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +HIDDEN_SIZE = 6144 +NUM_EXPERTS = 256 +MAX_TOKENS = 16 + + +class GlmRouterGemm(OpWrapper): + """`hidden_states @ weight.T` in FP32 from BF16 operands: the router logits `[tokens, 256]`. + + `weight` is the gate weight as stored, `[256, 6144]` BF16 row-major. One CTA per expert. Before + its grid-dependency wait each CTA prefetches its weight row into L2; right after the wait, + before any logit is written, it lets its dependents launch. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The c1 decode step routes the MTP layer's BF16 MoE here once per draft step (7 calls); the + # 75 MoE layers of the target compute their router logits inside the routing kernel. + CELLS: tuple[Cell, ...] = ( + Cell( + why="MTP draft steps: one row (draft step 0 after its row select, steps 1-6)", + spec=dict(tokens=1), + ), + Cell( + why="8 rows: a verify-width step routed without the row select", + spec=dict(tokens=8), + ), + Cell(why="16 tokens: the op's ceiling, its largest instance", spec=dict(tokens=16)), + ) + + note = """ + PDL (`TRTLLM_ENABLE_PDL`): the kernel reads only `weight` before its grid-dependency wait (an + L2 prefetch), so no kernel launched earlier in the step may write it. Its dependents launch + as soon as it passes the wait, before the logits exist: each must wait on its own grid + dependency before it reads the logits or writes `hidden_states`. At c1 the dependent is the + low-M BF16 MoE's FC1 kernel, which waits before it routes; its FC2 kernel reads the logits + before its own wait, which holds because FC1 releases it only after FC1's wait. + + Preconditions, enforced loudly by the op: CUDA BF16 2-D `hidden_states` and `weight` on one + device, contiguous and 16-byte aligned; hidden 6144; 256 experts; 1 to 16 tokens. + + Numerics: per thread an FP32 FMA chain over its 48 products, then a warp butterfly and the four + warp sums in order; deterministic, and bitwise the logits of the shared + `dsv3_router_gemm_op`, whose kernel sums in the same order. + """ + + def __call__(self, hidden_states: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + return torch.ops.trtllm.glm_router_gemm(hidden_states, weight) + + def reference(self, hidden_states: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """`hidden_states @ weight.T` accumulated in FP64 and rounded to FP32 once. + + The op's output is FP32, so an FP32-accumulated reference would carry rounding of the + same size as the kernel's; FP64 keeps the reference the more accurate side. + """ + return (hidden_states.double() @ weight.double().t()).float() + + def is_valid(self, hidden_states: torch.Tensor, weight: torch.Tensor) -> None: + # The shape gates of the kernel, which a target checks to route anything else to the + # shared dsv3_router_gemm_op. The op raises on them as well. + assert hidden_states.dim() == 2 and hidden_states.shape[1] == HIDDEN_SIZE, ( + f"hidden_states must be [tokens, {HIDDEN_SIZE}]" + ) + # One CTA per expert row of the [256, 6144] gate weight. + assert tuple(weight.shape) == (NUM_EXPERTS, HIDDEN_SIZE), ( + f"weight must be [{NUM_EXPERTS}, {HIDDEN_SIZE}], got {tuple(weight.shape)}" + ) + # One kernel instance per token count. + assert 1 <= hidden_states.shape[0] <= MAX_TOKENS, ( + f"the kernel takes 1 to {MAX_TOKENS} tokens, got {hidden_states.shape[0]}" + ) + + +router_gemm = GlmRouterGemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/lowm_bf16_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/lowm_bf16_moe.py new file mode 100644 index 000000000000..f2b83f9f4ede --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/lowm_bf16_moe.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One token through the MTP layer's BF16 MoE, routing included, via the glm_lowm_bf16_moe op.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + + +class GlmLowmBf16Moe(OpWrapper): + """The routed half of the MTP layer's MoE for one token, from the router logits to the weighted expert sum. + + Routing is the n_group == 1 case of `noaux_tc_op`: experts by `sigmoid(logits) + bias` (sigmoid in the tanh + form), the top 8 in descending order with the lower expert id first on ties, weights from the unbiased scores + normalized over the 8 and scaled by `routed_scaling_factor`. Then per routed expert `e_k`: + `act_k = bf16(silu(bf16(W1[e_k] x)) * bf16(W3[e_k] x))` and the output is + `bf16(sum_k bf16(w_k * (W2[e_k] act_k)))`, the sum in FP32 in the routing order. + + `w3_w1` is `[256, 2 * I, H]` with the up rows (W3) first and the gate rows (W1) second, `w2` is `[256, H, I]`, + all experts local (tensor-parallel slices, no expert parallelism). Returns `[1, H]`; the shared experts are not + part of the op. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The MTP layer (layer 78) is BF16 throughout; at tp16 an expert's intermediate slice is 2048 / 16 = 128. Its + # MoE runs on one row on every draft step (step 0 after the row select). The checkpoint's correction bias is + # FP32 and its routed scaling factor 2.5. The value cells drive the routing's three regimes. + CELLS: tuple[Cell, ...] = ( + Cell( + why="draft step at tp16: one token, hidden 6144, expert slice 128, random routing", + spec=dict(hidden=6144, inter=128, routing="random"), + ), + Cell( + why="tied routing keys at the top-8 boundary: the lower expert id wins", + spec=dict(hidden=6144, inter=128, routing="ties"), + ), + Cell( + why="selected scores summing below 2^-13: the weights take the FP64 expression", + spec=dict(hidden=6144, inter=128, routing="tiny"), + ), + ) + + note = """ + The kernels serve hidden sizes that are multiples of 2048 up to 8192 and expert slices that are multiples of + 64 up to 512, exactly 256 experts and top 8; the op refuses anything else, and more than one token, loudly. A + BF16 correction bias is accepted as well; the checkpoint's is FP32, so no cell drives it. + + The order of the two halves of `w3_w1` is not visible in the shapes: a gate-first tensor gives silently wrong + output. The target builds `w3_w1` up-first from the checkpoint's `up_proj` and `gate_proj`. + + Under PDL the FC2 kernel reads the router logits before its grid-dependency wait: the logits must be complete + when the FC1 kernel's wait returns, which holds for any producer earlier in stream order. + """ + + def __call__( + self, + x: torch.Tensor, + router_logits: torch.Tensor, + e_score_correction_bias: torch.Tensor, + w3_w1: torch.Tensor, + w2: torch.Tensor, + top_k: int, + routed_scaling_factor: float, + ) -> torch.Tensor: + return torch.ops.trtllm.glm_lowm_bf16_moe( + x, router_logits, e_score_correction_bias, w3_w1, w2, top_k, routed_scaling_factor + ) + + def reference( + self, + x: torch.Tensor, + router_logits: torch.Tensor, + e_score_correction_bias: torch.Tensor, + w3_w1: torch.Tensor, + w2: torch.Tensor, + top_k: int, + routed_scaling_factor: float, + ) -> torch.Tensor: + """Plain torch with the op's rounding points. + + BF16 after FC1, after SwiGLU and after each expert's scaled FC2 output, the routing weights in FP64 as + `noaux_tc_op` forms them, the expert sum in FP32 in the routing order. Products and sums are FP32 (no + TF32). + """ + inter = w2.shape[2] + scores = 0.5 * torch.tanh(0.5 * router_logits.float()[0]) + 0.5 + choice = scores + e_score_correction_bias.float() + ids = torch.sort(choice, descending=True, stable=True).indices[:top_k] + selected = scores[ids] + total = selected.sum().double() + weights = (selected.double() * routed_scaling_factor / (total + 1e-20)).float() + xf = x.float()[0] + out = torch.zeros(w2.shape[1], dtype=torch.float32, device=x.device) + with true_fp32_matmul(): + for k in range(top_k): + expert = int(ids[k]) + fc1 = w3_w1[expert].float() @ xf + up = fc1[:inter].to(torch.bfloat16).float() + gate = fc1[inter:].to(torch.bfloat16).float() + act = (gate * (1.0 / (1.0 + torch.exp(-gate))) * up).to(torch.bfloat16).float() + term = (weights[k] * (w2[expert].float() @ act)).to(torch.bfloat16).float() + out = out + term + return out.to(torch.bfloat16).unsqueeze(0) + + def is_valid( + self, + x: torch.Tensor, + router_logits: torch.Tensor, + e_score_correction_bias: torch.Tensor, + w3_w1: torch.Tensor, + w2: torch.Tensor, + top_k: int, + routed_scaling_factor: float, + ) -> None: + # Nothing to add: the op refuses every token count, expert count, top-k, shape, dtype, layout and + # alignment it does not serve, loudly, and the one silent mistake left -- a gate-first `w3_w1` -- is not + # visible in the tensors (see `note`). + del x, router_logits, e_score_correction_bias, w3_w1, w2, top_k, routed_scaling_factor + + +glm_lowm_bf16_moe = GlmLowmBf16Moe() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/routed_experts.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/routed_experts.py new file mode 100644 index 000000000000..6637fcbbe700 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/routed_experts.py @@ -0,0 +1,408 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GLM-5.2's routed experts at MoE TP16: DeepSeek-V3 routing without groups +(sigmoid + bias, top-8 of 256, scaled sum normalization), grouped NVFP4 FC1 + +SwiGLU + NVFP4 requantization + grouped NVFP4 FC2, in one trtllm-gen call. + +On the decode step (up to 8 rows) the call runs the router GEMV inside the +routing kernel, lets the NVFP4 input still be in flight on another stream +until FC1 starts, and runs FC2 on the K = 128 kernel; it stops before the +finalize step, which the post-MoE all-reduce hop performs.""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( # noqa: E501 + dequantize_nvfp4, + quantize_nvfp4, +) +from tensorrt_llm._torch.custom_ops.trtllm_gen_custom_ops import ( + FUSED_ROUTER_WORKSPACE_WORDS, + make_fused_router_workspace, +) + +__all__ = [ + "ACT_TYPE_SWIGLU", + "FUSED_ROUTER_MAX_TOKENS", + "FUSED_ROUTER_WORKSPACE_WORDS", + "ROUTING_METHOD_DEEPSEEK_V3", + "make_fused_router_workspace", + "routed_experts", +] + +#: `routing_method_type` of DeepSeek-V3 routing; with `n_group` 1 the op runs +#: it as sigmoid + bias, top-k, and the sum-normalized, scaled sigmoid weights. +ROUTING_METHOD_DEEPSEEK_V3 = 2 + +#: `act_type` for SwiGlu. +ACT_TYPE_SWIGLU = 0 + +#: The largest row count the fused router + routing kernel takes. +FUSED_ROUTER_MAX_TOKENS = 8 + +#: One expert stack's `(codes [E, N, K] uint8, block scales [E, N, K/16] +#: float8_e4m3fn)`, unshuffled and unswizzled. +Operand = Tuple[torch.Tensor, torch.Tensor] + +_IDLE_BITS = -(2**31) # -0.0f, the idle value of a fused-router exchange buffer + + +def _dequantize_activation(hidden_states: torch.Tensor, hidden_states_scale: torch.Tensor): + """`[T, H]` fp32 value of an NVFP4 activation in the op's input layout: + packed e2m1 codes `[T, H/2]` (element 2i in the low nibble) and linear e4m3 + block scales, one per 16 elements of a row. The global scale stays in.""" + rows = hidden_states.shape[0] + codes = torch.stack([hidden_states & 0x0F, hidden_states >> 4], dim=-1).reshape(rows, -1) + scales = hidden_states_scale.view(torch.float8_e4m3fn).reshape(rows, -1) + return dequantize_nvfp4(codes, scales) + + +class GlmRoutedExperts(OpWrapper): + """One routed-expert MoE layer of GLM-5.2 at MoE TP16, over NVFP4 input. + + Routing happens inside the op, from one of two sources: + + * `routing_logits` `[T, 256]` -- the router GEMM's fp32 logits; + * `router_hidden_states` `[T, 6144]` bf16 with `router_weight` + `[256, 6144]` bf16 and `router_workspace` (`make_fused_router_workspace`): + the routing kernel runs the router GEMV itself, bit for bit the logits of + `dsv3_router_gemm_op`, at 1-8 rows. `routing_logits` then only gives the + shape and is not read. + + `hidden_states_ready_event` is the raw `cudaEvent_t` of a `torch.cuda.Event` + recorded on the stream that writes `hidden_states` / `hidden_states_scale` + (`event.cuda_event`, after its first record): routing launches without + waiting for it, FC1 waits. `prefer_gemm2_k128` runs FC2 on the persistent + K = 128 kernel where it covers the call -- intermediate 128, tile N 8 and + enough SMs, which an 8-row call meets on Rubin's SM count -- with the Gemm2 + cubin's rows, bitwise. + + `do_finalize=False` returns `[fc2 rows [P, H] bf16, expert weights [T, 8] + bf16, expanded index -> permuted row [T, 8] int32]`: token `t`'s output is + `sum_k weights[t, k] * rows[map[t, k]]`. `do_finalize=True` returns + `[output [T, H] bf16]`. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # GLM-5.2's routed geometry at MoE TP16, ep 1: every rank holds all 256 + # experts with a 2048 / 16 = 128-wide intermediate slice. + _GLM_TP16 = dict( + hidden=6144, + intermediate=128, + num_experts=256, + top_k=8, + routed_scaling_factor=2.5, + local_expert_offset=0, + local_num_experts=256, + ) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="the c1 verify step: 8 rows routed by the fused router GEMV, the input joined " + "right before FC1, FC2 at K = 128, unfinalized for the post-MoE hop", + spec=dict(tokens=8, fused_router=True, late_join=True, gemm2_k128=True, **_GLM_TP16), + ), + Cell( + why="one row, a step without MTP: the fused router's single-token instance and " + "FC2 at K = 128 over 8 token tiles", + spec=dict(tokens=1, fused_router=True, late_join=True, gemm2_k128=True, **_GLM_TP16), + ), + Cell( + why="32 rows (4 requests x 8), past the fused router: routed from the router " + "GEMM's logits, FC2 back on the Gemm2 cubin", + spec=dict(tokens=32, fused_router=False, late_join=False, gemm2_k128=True, **_GLM_TP16), + ), + ) + + note = """ + Certified unfinalized (`do_finalize=False`) only; the finalize step is the + post-MoE hop's. No expert-parallel window (`local_expert_offset` 0, all 256 + experts local), no `output` buffer, `use_dp` False, `topk_ids` / + `topk_weights` None. `gemm1_bias`, `gemm1_alpha`, `gemm1_beta`, + `gemm1_clamp_limit` and `gemm2_bias` are None, `act_type` SwiGlu. + + `routing_bias` is GLM-5.2's `e_score_correction_bias` in bf16, the dtype + the TRTLLM MoE backend loads it in; `n_group` / `topk_group` are 1. + + The tactic is the autotuner's: the cells run cold (the op's default tile, + N 8 at these row counts). Under autotuning, a call routed by the fused + router profiles from `routing_logits`, so the caller fills that buffer + once (random logits) rather than leaving it uninitialized. + + A router workspace is per call site and its launches must run in stream + order; the kernel leaves the exchange buffer of the next launch idle and + advances the epoch word. `tune_max_num_tokens` is an autotuner bucket hint. + + The caller owns the weight layout -- FC1 rows `[up ; gate]` interleaved and + 32-row block shuffled, FC2 rows block shuffled, both scale stacks 128x4 + swizzled -- as the catalog `fp4_block_scale_moe_runner` entry describes. + """ + + def __call__( + self, + routing_logits: torch.Tensor, + routing_bias: torch.Tensor, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm1_bias: Optional[torch.Tensor], + gemm1_alpha: Optional[torch.Tensor], + gemm1_beta: Optional[torch.Tensor], + gemm1_clamp_limit: Optional[torch.Tensor], + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + gemm2_bias: Optional[torch.Tensor], + output1_scale_scalar: torch.Tensor, + output1_scale_gate_scalar: torch.Tensor, + output2_scale_scalar: torch.Tensor, + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: float, + routing_method_type: int, + do_finalize: bool, + act_type: int, + topk_weights: Optional[torch.Tensor], + topk_ids: Optional[torch.Tensor], + output: Optional[torch.Tensor], + tune_max_num_tokens: int, + use_dp: bool, + hidden_states_ready_event: Optional[int], + router_hidden_states: Optional[torch.Tensor], + router_weight: Optional[torch.Tensor], + router_workspace: Optional[torch.Tensor], + prefer_gemm2_k128: bool, + ) -> list[torch.Tensor]: + return torch.ops.trtllm.fp4_block_scale_moe_runner( + routing_logits, + routing_bias, + hidden_states, + hidden_states_scale, + gemm1_weights, + gemm1_weights_scale, + gemm1_bias, + gemm1_alpha, + gemm1_beta, + gemm1_clamp_limit, + gemm2_weights, + gemm2_weights_scale, + gemm2_bias, + output1_scale_scalar, + output1_scale_gate_scalar, + output2_scale_scalar, + num_experts, + top_k, + n_group, + topk_group, + intermediate_size, + local_expert_offset, + local_num_experts, + routed_scaling_factor, + routing_method_type, + do_finalize, + act_type, + topk_weights=topk_weights, + topk_ids=topk_ids, + output=output, + tune_max_num_tokens=tune_max_num_tokens, + use_dp=use_dp, + hidden_states_ready_event=hidden_states_ready_event, + router_hidden_states=router_hidden_states, + router_weight=router_weight, + router_workspace=router_workspace, + prefer_gemm2_k128=prefer_gemm2_k128, + ) + + def reference( + self, + routing_logits: Optional[torch.Tensor], + routing_bias: torch.Tensor, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + up: Operand, + gate: Operand, + down: Operand, + output1_scale_scalar: torch.Tensor, + output1_scale_gate_scalar: torch.Tensor, + output2_scale_scalar: torch.Tensor, + top_k: int, + local_expert_offset: int, + routed_scaling_factor: float, + router_hidden_states: Optional[torch.Tensor] = None, + router_weight: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """The layer, in plain torch over dequantized weights, fp32 throughout. + + Returns `(output [T, H] bf16, topk ids [T, top_k], topk weights [T, + top_k] bf16)`: the finalized layer, rounded once, and the routing it + used. A `do_finalize=False` call is compared after summing its rows with + its own weights. + + The documented exception to mirroring: `__call__` takes the prepared + weight operands -- shuffled rows, swizzled scales -- and a reference + given those would re-derive the layout rather than the computation, so + this takes the per-expert `(codes, block scales)` pairs `up`, `gate` + and `down` the layout was built from. Everything else keeps the name it + has in `__call__`. + + Routing: the logits are `routing_logits`, or the router GEMV + `router_hidden_states @ router_weight.T` when those are given (the + fused router). Scores are `sigmoid(logits)` -- in the routing kernel's + `0.5 tanh(x / 2) + 0.5` form --, experts are chosen by score + bias, + and an expert's weight is its score over the sum of the chosen scores, + times `routed_scaling_factor`, rounded to bf16 as the kernel stores it. + + Experts: the op's own scalars carry the global scales. FC1's gate half + is dequantized by `output1_scale_gate_scalar` (`a1`), its linear half + by the same `a1`; the SwiGLU output is requantized to NVFP4 under the + FC2-input global scale `g2 = output1_scale_scalar / a1`, which the FC1 + epilogue folds in, and FC2's accumulator is dequantized by + `output2_scale_scalar`. Skipping that requantization is a reference + some 20 ulp RMS away from a correct kernel. + """ + x = _dequantize_activation(hidden_states, hidden_states_scale) + num_tokens, hidden = x.shape + with true_fp32_matmul(): + if router_hidden_states is not None: + logits = router_hidden_states.float() @ router_weight.float().t() + else: + logits = routing_logits.float() + scores = 0.5 * torch.tanh(0.5 * logits) + 0.5 + chosen = torch.topk(scores + routing_bias.float(), top_k, dim=-1).indices + picked = scores.gather(1, chosen) + weights = (picked * routed_scaling_factor / picked.sum(-1, keepdim=True)).to( + torch.bfloat16 + ) + + up_c, up_s = up + gt_c, gt_s = gate + dn_c, dn_s = down + out = torch.zeros(num_tokens, hidden, dtype=torch.float32, device=x.device) + for local_e in range(up_c.shape[0]): + tok, slot = (chosen == local_expert_offset + local_e).nonzero(as_tuple=True) + if tok.numel() == 0: + continue + a1 = output1_scale_gate_scalar[local_e].float() + g2 = output1_scale_scalar[local_e].float() / a1 + a2 = output2_scale_scalar[local_e].float() + xe = x[tok] + u = (xe @ dequantize_nvfp4(up_c[local_e], up_s[local_e]).t()) * a1 + g = (xe @ dequantize_nvfp4(gt_c[local_e], gt_s[local_e]).t()) * a1 + act = quantize_nvfp4(u * g * torch.sigmoid(g), float(g2)) + y = (act @ dequantize_nvfp4(dn_c[local_e], dn_s[local_e]).t()) * a2 + out.index_add_(0, tok, y * weights[tok, slot].float().unsqueeze(1)) + return out.to(torch.bfloat16), chosen.to(torch.int32), weights + + def is_valid( + self, + routing_logits: torch.Tensor, + routing_bias: torch.Tensor, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm1_bias: Optional[torch.Tensor], + gemm1_alpha: Optional[torch.Tensor], + gemm1_beta: Optional[torch.Tensor], + gemm1_clamp_limit: Optional[torch.Tensor], + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + gemm2_bias: Optional[torch.Tensor], + output1_scale_scalar: torch.Tensor, + output1_scale_gate_scalar: torch.Tensor, + output2_scale_scalar: torch.Tensor, + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: float, + routing_method_type: int, + do_finalize: bool, + act_type: int, + topk_weights: Optional[torch.Tensor], + topk_ids: Optional[torch.Tensor], + output: Optional[torch.Tensor], + tune_max_num_tokens: int, + use_dp: bool, + hidden_states_ready_event: Optional[int], + router_hidden_states: Optional[torch.Tensor], + router_weight: Optional[torch.Tensor], + router_workspace: Optional[torch.Tensor], + prefer_gemm2_k128: bool, + ) -> None: + """The ways this op takes an input and answers it wrongly. + + Shapes, dtypes, the routing configuration the fused router covers and + the workspace size are rejected loudly by the op and are not restated. + """ + # The kernels take raw data pointers and assume dense row-major + # tensors; a strided view is read as if dense. + for name, tensor in ( + ("routing_logits", routing_logits), + ("routing_bias", routing_bias), + ("hidden_states", hidden_states), + ("hidden_states_scale", hidden_states_scale), + ("gemm1_weights", gemm1_weights), + ("gemm1_weights_scale", gemm1_weights_scale), + ("gemm2_weights", gemm2_weights), + ("gemm2_weights_scale", gemm2_weights_scale), + ("output1_scale_scalar", output1_scale_scalar), + ("output1_scale_gate_scalar", output1_scale_gate_scalar), + ("output2_scale_scalar", output2_scale_scalar), + ("router_hidden_states", router_hidden_states), + ("router_weight", router_weight), + ("router_workspace", router_workspace), + ("output", output), + ): + if tensor is not None: + assert tensor.is_contiguous(), ( + f"{name} must be contiguous; a strided view is read as if dense " + "and silently produces wrong results" + ) + # The prepared block-scale layout is misread at a hidden size that is + # only a multiple of 128 (see the catalog fp4_block_scale_moe_runner). + assert gemm1_weights.shape[-1] * 2 % 256 == 0, ( + "hidden size (gemm1_weights.shape[-1] * 2) must be a multiple of 256; " + "the prepared block-scale layout is silently misread otherwise" + ) + # A torch.cuda.Event has no cudaEvent_t before its first record, and a + # null handle reads as "no event": FC1 would not wait for the input. + assert hidden_states_ready_event is None or hidden_states_ready_event != 0, ( + "hidden_states_ready_event is a null cudaEvent_t: take event.cuda_event after " + "the event is recorded, or FC1 reads hidden_states without waiting" + ) + # The routing CTA takes any value but -0.0 in the exchange buffer of + # this launch as a published slice sum. Unless that buffer is idle, it + # routes on stale sums -- a zero-filled workspace routes on zero logits. + if router_hidden_states is not None and router_workspace is not None: + words = router_workspace.view(torch.int32) + buffer_words = (FUSED_ROUTER_WORKSPACE_WORDS - 4) // 2 + epoch = int(words[2 * buffer_words].item()) + current = words[(epoch & 1) * buffer_words : ((epoch & 1) + 1) * buffer_words] + assert bool((current == _IDLE_BITS).all()), ( + "router_workspace is not in its between-launch state (the exchange buffer " + "of the next launch must hold -0.0 everywhere); build it with " + "make_fused_router_workspace and let only this op write it" + ) + + +routed_experts = GlmRoutedExperts() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/norm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/norm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_exchange.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_exchange.py new file mode 100644 index 000000000000..cdc715a2f082 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_exchange.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MNNVL Lamport workspace the two draft-step ops exchange through, and their shape constants. + +``glm_draft_sample_glue`` and ``glm_draft_ehproj_table`` publish each rank's (index, value) packet (and, for the +table op, a token-table row and the rank's hidden-half outputs) through the multicast address of a dedicated +workspace and poll the local copy until every rank's slot has landed. The workspace has three Lamport buffers and +the nine 32-bit flags of the MNNVL all-reduce layout (current index, dirty index, bytes per buffer, dirty stages, +four byte counts to clear, an access counter); an unwritten word holds 0xFFFFFFFF rather than the all-reduce's -0.0. + +Each op empties the buffer of its previous call with the threads of its current launch, so a workspace serves one +op and nothing else. +""" + +from __future__ import annotations + +import os +from typing import Any + +import torch + +from tensorrt_llm._torch.distributed.ops import _get_mnnvl_workspace_comm, _mnnvl_device_index +from tensorrt_llm._utils import torch_pybind11_abi +from tensorrt_llm.bindings.internal.runtime import McastGPUBuffer +from tensorrt_llm.mapping import Mapping + +#: Rows (requests) one draft-step launch serves (kDraftSampleGlueMaxRows). +MAX_ROWS = 8 +#: Bytes of one exchanged (index, value) packet (kDraftSampleGluePacketBytes). +PACKET_BYTES = 16 +#: Hidden size the draft-step kernels are built for (kDraftSampleGlueHidden). +HIDDEN_SIZE = 6144 +#: Largest TP size of glm_draft_ehproj_table (kDraftEhprojTableMaxRanks). +EHPROJ_TABLE_MAX_RANKS = 16 + + +def ehproj_table_supports_tp(tp_size: int) -> bool: + """Whether glm_draft_ehproj_table serves ``tp_size`` ranks (draftEhprojTableSupportsTpSize). + + 2 to 16 ranks whose shares of the hidden size are whole 8-column vectors. + """ + return 2 <= tp_size <= EHPROJ_TABLE_MAX_RANKS and HIDDEN_SIZE % (8 * tp_size) == 0 + + +def sample_glue_slot_bytes(tp_size: int) -> int: + """Bytes one rank publishes per row in glm_draft_sample_glue: its (index, value) packet.""" + del tp_size + return PACKET_BYTES + + +def ehproj_table_slot_bytes(tp_size: int) -> int: + """Bytes one rank publishes per row in glm_draft_ehproj_table (draftEhprojTableSlotBytes). + + The packet, the BF16 token-table row of its candidate and its FP32 share of the hidden-half outputs. + """ + return PACKET_BYTES + 2 * HIDDEN_SIZE + 4 * HIDDEN_SIZE // tp_size + + +def _all_succeeded(comm: Any, ok: bool) -> bool: + """Whether ``ok`` holds on every rank of ``comm`` (an MPI communicator or a ProcessGroup with a CPU backend).""" + if isinstance(comm, torch.distributed.ProcessGroup): + flag = torch.tensor([1 if ok else 0], dtype=torch.int32) + work = comm._get_backend(torch.device("cpu")).allreduce([flag]) + work.wait() + return int(flag.item()) == comm.size() + return comm.allreduce(int(ok)) == comm.Get_size() + + +class DraftExchangeWorkspace: + """One draft-step op's MNNVL workspace over a TP group. + + ``buffer`` is the local unicast view of the three Lamport buffers (int32, every word 0xFFFFFFFF) and ``flags`` + the nine uint32 Lamport flags. Construction is collective over the group and either succeeds on every rank or + raises on every rank. Build it outside CUDA-graph capture, once per op; the captured graphs keep its addresses. + """ + + NUM_LAMPORT_BUFFERS = 3 + + def __init__( + self, + buffer_size_bytes: int, + group_size: int, + group_rank: int, + device_index: int, + use_fabric_handle: bool, + comm: Any, + ): + """Allocate ``3 * buffer_size_bytes`` of multicast memory over the group and reset it. + + ``comm`` is the group's MPI communicator or a ``torch.distributed.ProcessGroup`` with a CPU backend, ordered + by group rank. + """ + assert buffer_size_bytes > 0 and buffer_size_bytes % 16 == 0, buffer_size_bytes + self.buffer_size_bytes = buffer_size_bytes + self.group_size = group_size + workspace_size_bytes = self.NUM_LAMPORT_BUFFERS * buffer_size_bytes + error = None + state = None + try: + if isinstance(comm, torch.distributed.ProcessGroup): + handle = McastGPUBuffer( + workspace_size_bytes, + group_size, + group_rank, + device_index, + use_fabric_handle, + process_group=comm, + pybind11_abi=torch_pybind11_abi(), + ) + else: + handle = McastGPUBuffer( + workspace_size_bytes, + group_size, + group_rank, + device_index, + use_fabric_handle, + comm.py2f(), + ) + buffer = handle.get_uc_buffer(group_rank, (workspace_size_bytes // 4,), torch.int32, 0) + with torch.inference_mode(): + buffer.fill_(-1) + flags = torch.tensor( + [0, 2, buffer_size_bytes, 0, 0, 0, 0, 0, 0], + dtype=torch.uint32, + device=torch.device("cuda", device_index), + ) + torch.cuda.synchronize() + state = (handle, buffer, flags) + except (RuntimeError, ValueError) as e: + error = e + if not _all_succeeded(comm, error is None): + raise RuntimeError( + "draft-step MNNVL workspace construction failed on at least one rank" + ) from error + # The handle keeps the multicast allocation alive; the communicator outlives it. + self._handle, self.buffer, self.flags = state + self._comm = comm + + @classmethod + def for_mapping(cls, mapping: Mapping, buffer_size_bytes: int) -> "DraftExchangeWorkspace": + """The workspace over ``mapping``'s TP group, with the MNNVL all-reduce workspace's communicator. + + Fabric handles across nodes or under ``TRTLLM_FORCE_MNNVL_AR``, NVLS handles within a node, as the MNNVL + all-reduce workspace takes them. + """ + use_fabric_handle = ( + os.environ.get("TRTLLM_FORCE_MNNVL_AR", "0") == "1" or mapping.is_multi_node() + ) + return cls( + buffer_size_bytes, + mapping.tp_size, + mapping.tp_rank, + _mnnvl_device_index(mapping), + use_fabric_handle, + _get_mnnvl_workspace_comm(mapping), + ) + + @classmethod + def for_sample_glue(cls, mapping: Mapping) -> "DraftExchangeWorkspace": + """The workspace of glm_draft_sample_glue for up to MAX_ROWS rows.""" + slot_bytes = sample_glue_slot_bytes(mapping.tp_size) + return cls.for_mapping(mapping, MAX_ROWS * mapping.tp_size * slot_bytes) + + @classmethod + def for_ehproj_table(cls, mapping: Mapping) -> "DraftExchangeWorkspace": + """The workspace of glm_draft_ehproj_table for up to MAX_ROWS rows.""" + slot_bytes = ehproj_table_slot_bytes(mapping.tp_size) + return cls.for_mapping(mapping, MAX_ROWS * mapping.tp_size * slot_bytes) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_glue_rmsnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_glue_rmsnorm.py new file mode 100644 index 000000000000..6f4a3ba08625 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_glue_rmsnorm.py @@ -0,0 +1,457 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CuTe DSL kernel behind ``draft_glue_rmsnorm_cat``: the MTP draft-step eh_proj input in one kernel. + +The MTP layer feeds ``eh_proj`` with this rank's columns of ``cat([enorm(embed(input_ids)), hnorm(hidden_states)], +-1)``. ``DraftGlueRMSNormKernel`` gathers the embedding rows from the table itself (a zero row for an id outside the +table), normalizes them with the enorm weight and the hidden rows with the hnorm weight, and stores only the +requested window of the concatenated rows, contiguously. + +Each row is normalized as FlashInfer's CuTe ``RMSNormKernel`` (the kernel ``flashinfer.norm.rmsnorm`` runs for 2D +BF16 inputs) does it: the thread geometry is taken from an ``RMSNormKernel`` instance (64 threads per row for +H = 6144, 8-element vectors, the same column order per thread), the reduction is FlashInfer's +``row_reduce_sum_multirow``, the arithmetic statements are FlashInfer's, and the kernel is compiled with the same +options. The per-row code follows FlashInfer's ``QKRMSNormKernel``, which gives every row its own tile so that a +row can come from any tensor. +""" + +from __future__ import annotations + +import functools + +import cutlass +import cutlass.cute as cute +import torch +from cutlass import Float32, Int64 +from flashinfer.norm.kernels.rmsnorm import RMSNormKernel +from flashinfer.norm.utils import make_tv_layout, predicate_k, row_reduce_sum_multirow + +_IDS_DTYPES = {torch.int32: cutlass.Int32, torch.int64: cutlass.Int64} + + +@cute.jit +def _predicate_cols( + tXcX: cute.Tensor, lo: cutlass.Constexpr[int], hi: cutlass.Constexpr[int] +) -> cute.Tensor: + """Copy predicate that is true where a vector's first column lies in [lo, hi).""" + tXpX = cute.make_rmem_tensor( + cute.make_layout( + ( + cute.size(tXcX, mode=[0, 1]), + cute.size(tXcX, mode=[1]), + cute.size(tXcX, mode=[2]), + ), + stride=(cute.size(tXcX, mode=[2]), 0, 1), + ), + cutlass.Boolean, + ) + for rest_v in cutlass.range_constexpr(tXpX.shape[0]): + for rest_k in cutlass.range_constexpr(tXpX.shape[2]): + col = tXcX[(0, rest_v), 0, rest_k][1] + in_window = cute.elem_less(col, hi) + if cutlass.const_expr(lo > 0): + in_window = in_window & cute.elem_less(lo - 1, col) + tXpX[rest_v, 0, rest_k] = in_window + return tXpX + + +class DraftGlueRMSNormKernel: + """RMSNorm of the embedding rows (half 0) and of the hidden rows (half 1) of the MTP draft step. + + Column ``k`` of half ``h`` is column ``h * H + k`` of the concatenated row. The kernel computes the halves + overlapping the column window ``[col_lo, col_hi)`` and stores the window's columns at + ``out[row, col - col_lo]``. Half 0 reads ``table[input_ids[row]]`` (a zero row for an id outside the table). Each + group of ``threads_per_row`` threads normalizes one (row, half) pair; with both halves a CTA covers one row. + + Args: + hidden_size: H, the length of each normalized row. + half_begin: First half overlapping the window. + num_halves: Number of halves overlapping the window (1 or 2). + col_lo: First concatenated column stored. + col_hi: One past the last concatenated column stored. + """ + + def __init__( + self, + hidden_size: int, + half_begin: int, + num_halves: int, + col_lo: int, + col_hi: int, + ): + reference = RMSNormKernel(cutlass.BFloat16, hidden_size, weight_bias=0.0) + if reference.cluster_n != 1: + raise ValueError( + f"hidden size {hidden_size} needs a CTA cluster per row in RMSNormKernel; " + "the draft glue kernel covers single-CTA rows only" + ) + self.H = hidden_size + self.weight_bias = reference.weight_bias + self.copy_bits = reference.copy_bits + self.vec_size = reference.vec_size + self.threads_per_row = reference.threads_per_row + self.num_threads = reference.num_threads + self.rows_per_block = reference.rows_per_block + self.warps_per_row = reference.warps_per_row + self.num_vec_blocks = reference.num_vec_blocks + self.cols_per_tile = reference.cols_per_tile + self.use_async_copy = reference.use_async_copy + + self.half_begin = half_begin + self.num_halves = num_halves + self.out_cols = col_hi - col_lo + self.out_col_shift = col_lo + # Stored columns [lo, hi) of each half, in that half's column coordinates. + windows = [] + for h in range(2): + lo = min(max(col_lo - h * hidden_size, 0), hidden_size) + hi = max(min(col_hi - h * hidden_size, hidden_size), lo) + windows.append((lo, hi)) + self.windows = tuple(windows) + + def _smem_size_in_bytes(self) -> int: + tile_bytes = 0 + if self.use_async_copy: + tile_bytes = self.rows_per_block * self.cols_per_tile * (cutlass.BFloat16.width // 8) + return tile_bytes + self.rows_per_block * self.warps_per_row * 4 + + @cute.jit + def __call__( + self, + mSrc0: cute.Tensor, + mIds: cute.Tensor, + mSrc1: cute.Tensor, + mW0: cute.Tensor, + mW1: cute.Tensor, + mOut: cute.Tensor, + M: Int64, + vocab: Int64, + eps: Float32, + enable_pdl: cutlass.Constexpr[bool], + stream, + ): + tv_shape, tv_stride = make_tv_layout( + self.threads_per_row, self.vec_size, self.num_vec_blocks + ) + tv_layout = cute.make_layout(tv_shape, stride=tv_stride) + tiler_1d = (1, self.cols_per_tile) + + self.kernel( + mSrc0, mIds, mSrc1, mW0, mW1, mOut, M, vocab, eps, enable_pdl, tv_layout, tiler_1d + ).launch( + grid=[cute.ceil_div(M * self.num_halves, self.rows_per_block), 1, 1], + block=[self.num_threads, 1, 1], + smem=self._smem_size_in_bytes(), + stream=stream, + use_pdl=enable_pdl, + ) + + @cute.kernel + def kernel( + self, + mSrc0: cute.Tensor, + mIds: cute.Tensor, + mSrc1: cute.Tensor, + mW0: cute.Tensor, + mW1: cute.Tensor, + mOut: cute.Tensor, + M: Int64, + vocab: Int64, + eps: Float32, + enable_pdl: cutlass.Constexpr[bool], + tv_layout: cute.Layout, + tiler_1d: cute.Shape, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + + # PDL: the ids and the hidden rows may come from the preceding kernel, so wait before any load. The + # dependent (eh_proj) may launch right away; it waits before it reads this kernel's output. + if enable_pdl: + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + + H = self.H + weight_bias = self.weight_bias + copy_bits = self.copy_bits + threads_per_row = self.threads_per_row + rows_per_block = self.rows_per_block + warps_per_row = self.warps_per_row + num_halves = self.num_halves + + lane_in_row = tidx % threads_per_row + row_in_block = tidx // threads_per_row + slot = bidx * rows_per_block + row_in_block + if cutlass.const_expr(num_halves == 1): + row = slot + is_half0 = self.half_begin == 0 + else: + row = slot // num_halves + is_half0 = slot % num_halves == 0 + row_in_bounds = row < M + + # ===== Allocate shared memory ===== + smem = cutlass.utils.SmemAllocator() + + if cutlass.const_expr(self.use_async_copy): + sX = smem.allocate_tensor( + mSrc1.element_type, + cute.make_ordered_layout((rows_per_block, self.cols_per_tile), order=(1, 0)), + byte_alignment=16, + ) + + reduction_buffer = smem.allocate_tensor( + Float32, + cute.make_layout((rows_per_block, warps_per_row)), + byte_alignment=4, + ) + mbar_ptr = None + cluster_n = 1 + + # ===== Source rows ===== + # Half 0 reads the embedding row of the token, half 1 the hidden row. + valid0 = row_in_bounds + src_row0 = row + if cutlass.const_expr(self.half_begin == 0): + # Same guard as the Embedding module: an id outside [0, vocab), such as the -1 padding of rejected draft + # tokens, gives a zero embedding row. + token = Int64(-1) + if row_in_bounds: + token = Int64(mIds[row]) + valid0 = (token >= Int64(0)) & (token < vocab) + src_row0 = cutlass.select_(valid0, token, Int64(0)) + gX0 = cute.local_tile(mSrc0, tiler_1d, (src_row0, 0)) + gX1 = cute.local_tile(mSrc1, tiler_1d, (row, 0)) + + # ===== Output: the window columns of each half ===== + # Column k of half h is stored at out column h * H + k - out_col_shift. + out_layout = cute.make_layout((M, H), stride=(self.out_cols, 1)) + mO0 = cute.make_tensor(mOut.iterator + (0 - self.out_col_shift), out_layout) + mO1 = cute.make_tensor(mOut.iterator + (H - self.out_col_shift), out_layout) + gO0 = cute.local_tile(mO0, tiler_1d, (row, 0)) + gO1 = cute.local_tile(mO1, tiler_1d, (row, 0)) + + mW0_2d = cute.prepend_ones(mW0, up_to_rank=2) + mW1_2d = cute.prepend_ones(mW1, up_to_rank=2) + + # ===== Create TiledCopy atoms ===== + copy_atom_sync = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + mSrc1.element_type, + num_bits_per_copy=copy_bits, + ) + copy_atom_store = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + mOut.element_type, + num_bits_per_copy=copy_bits, + ) + + if cutlass.const_expr(self.use_async_copy): + copy_atom_async = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(), + mSrc1.element_type, + num_bits_per_copy=copy_bits, + ) + tiled_copy_load = cute.make_tiled_copy(copy_atom_async, tv_layout, tiler_1d) + else: + tiled_copy_load = cute.make_tiled_copy(copy_atom_sync, tv_layout, tiler_1d) + + tiled_copy_W = cute.make_tiled_copy(copy_atom_sync, tv_layout, tiler_1d) + tiled_copy_store = cute.make_tiled_copy(copy_atom_store, tv_layout, tiler_1d) + + thr_copy_X = tiled_copy_load.get_slice(lane_in_row) + thr_copy_W = tiled_copy_W.get_slice(lane_in_row) + thr_copy_O = tiled_copy_store.get_slice(lane_in_row) + + # ===== Partition input ===== + tXgX0 = thr_copy_X.partition_S(gX0) + tXgX1 = thr_copy_X.partition_S(gX1) + tXrX = cute.make_fragment_like(tXgX1) + + if cutlass.const_expr(self.use_async_copy): + sX_row = cute.local_tile(sX, tiler_1d, (row_in_block, 0)) + tXsX = thr_copy_X.partition_D(sX_row) + + # ===== Partition weight (sync, separate tiled copy) ===== + tWgW0 = thr_copy_W.partition_S(mW0_2d) + tWgW1 = thr_copy_W.partition_S(mW1_2d) + tWrW = cute.make_fragment_like(tWgW0) + tXrW = thr_copy_X.retile(tWrW) + + # ===== Partition output ===== + tXgO0 = thr_copy_O.partition_D(gO0) + tXgO1 = thr_copy_O.partition_D(gO1) + tXrO = cute.make_fragment_like(tXgO0) + + # ===== Bounds checking ===== + id1d = cute.make_identity_tensor(tiler_1d) + tXcX = thr_copy_X.partition_S(id1d) + tXpX = predicate_k(tXcX, limit=H) + tWpW = predicate_k(thr_copy_W.partition_S(id1d), limit=H) + lo0, hi0 = self.windows[0] + lo1, hi1 = self.windows[1] + if cutlass.const_expr(lo0 == 0 and hi0 == H): + tXpO0 = tXpX + else: + tXpO0 = _predicate_cols(tXcX, lo0, hi0) + if cutlass.const_expr(lo1 == 0 and hi1 == H): + tXpO1 = tXpX + else: + tXpO1 = _predicate_cols(tXcX, lo1, hi1) + + # ===== Pass 1: Load input + compute sum of squares ===== + if cutlass.const_expr(self.use_async_copy): + if is_half0: + if valid0: + cute.copy(copy_atom_async, tXgX0, tXsX, pred=tXpX) + else: + tXrX.store(cute.zeros_like(tXrX, dtype=mSrc1.element_type)) + cute.autovec_copy(tXrX, tXsX) + else: + if row_in_bounds: + cute.copy(copy_atom_async, tXgX1, tXsX, pred=tXpX) + cute.arch.cp_async_commit_group() + + if is_half0: + cute.copy(copy_atom_sync, tWgW0, tWrW, pred=tWpW) + else: + cute.copy(copy_atom_sync, tWgW1, tWrW, pred=tWpW) + + cute.arch.cp_async_wait_group(0) + + cute.autovec_copy(tXsX, tXrX) + else: + tXrX.store(cute.zeros_like(tXrX, dtype=mSrc1.element_type)) + if is_half0: + if valid0: + cute.copy(copy_atom_sync, tXgX0, tXrX, pred=tXpX) + cute.copy(copy_atom_sync, tWgW0, tWrW, pred=tWpW) + else: + if row_in_bounds: + cute.copy(copy_atom_sync, tXgX1, tXrX, pred=tXpX) + cute.copy(copy_atom_sync, tWgW1, tWrW, pred=tWpW) + + x = tXrX.load().to(Float32) + x_sq = x * x + sum_sq = row_reduce_sum_multirow( + x_sq, threads_per_row, reduction_buffer, mbar_ptr, cluster_n + ) + + mean_sq = sum_sq / Float32(H) + rstd = cute.math.rsqrt(mean_sq + eps, fastmath=True) + + cute.arch.barrier() + + # ===== Pass 2: Normalize and store output ===== + # Re-load x from shared memory to relieve register pressure. + if cutlass.const_expr(self.use_async_copy): + cute.autovec_copy(tXsX, tXrX) + x = tXrX.load().to(Float32) + + w = tXrW.load().to(Float32) + y = x * rstd * (w + Float32(weight_bias)) + + tXrO.store(y.to(mOut.element_type)) + + if row_in_bounds: + if is_half0: + cute.copy(copy_atom_store, tXrO, tXgO0, pred=tXpO0) + else: + cute.copy(copy_atom_store, tXrO, tXgO1, pred=tXpO1) + + +@functools.cache +def _compile_draft_glue_kernel( + hidden_size: int, + half_begin: int, + num_halves: int, + ids_dtype: torch.dtype, + col_lo: int, + col_hi: int, + enable_pdl: bool, +): + kernel_obj = DraftGlueRMSNormKernel(hidden_size, half_begin, num_halves, col_lo, col_hi) + + sym_m = cute.sym_int(64) + src0_rows = cute.sym_int(64) if half_begin == 0 else sym_m + src0_fake = cute.runtime.make_fake_tensor( + cutlass.BFloat16, + (src0_rows, hidden_size), + (cute.sym_int64(divisibility=kernel_obj.vec_size), 1), + assumed_align=16, + ) + ids_fake = cute.runtime.make_fake_compact_tensor(_IDS_DTYPES[ids_dtype], (sym_m,)) + src1_fake = cute.runtime.make_fake_tensor( + cutlass.BFloat16, + (sym_m, hidden_size), + (cute.sym_int64(divisibility=kernel_obj.vec_size), 1), + assumed_align=16, + ) + w0_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, (hidden_size,), assumed_align=16 + ) + w1_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, (hidden_size,), assumed_align=16 + ) + out_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, (sym_m, col_hi - col_lo), stride_order=(1, 0), assumed_align=16 + ) + stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + + # Same compile options as FlashInfer's RMSNorm kernels. + return cute.compile( + kernel_obj, + src0_fake, + ids_fake, + src1_fake, + w0_fake, + w1_fake, + out_fake, + Int64(1), + Int64(1), + Float32(1e-6), + enable_pdl, + stream_fake, + options="--enable-tvm-ffi", + ) + + +def draft_glue_rmsnorm_cat( + embed_table: torch.Tensor, + input_ids: torch.Tensor, + hidden_states: torch.Tensor, + enorm_weight: torch.Tensor, + hnorm_weight: torch.Tensor, + eps: float, + col_lo: int, + col_hi: int, + enable_pdl: bool, +) -> torch.Tensor: + """Columns ``[col_lo, col_hi)`` of ``cat([rmsnorm(embed_table[input_ids]), rmsnorm(hidden_states)], -1)``. + + The embedding rows are normalized with ``enorm_weight``, the hidden rows with ``hnorm_weight``, both with + ``eps``; only the halves overlapping the window are computed. Returns the window as a contiguous + ``[M, col_hi - col_lo]`` BF16 tensor. + """ + num_rows, hidden_size = hidden_states.shape + half_begin = col_lo // hidden_size + num_halves = (col_hi - 1) // hidden_size - half_begin + 1 + # Without the embedding half the table is not read; any [M, H] tensor fills the argument. + embed_src = embed_table if half_begin == 0 else hidden_states + out = hidden_states.new_empty((num_rows, col_hi - col_lo)) + if num_rows > 0: + kernel = _compile_draft_glue_kernel( + hidden_size, half_begin, num_halves, input_ids.dtype, col_lo, col_hi, enable_pdl + ) + kernel( + embed_src, + input_ids, + hidden_states, + enorm_weight, + hnorm_weight, + out, + num_rows, + embed_src.shape[0], + eps, + ) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_step_state.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_step_state.py new file mode 100644 index 000000000000..65ba0b8f088d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_step_state.py @@ -0,0 +1,189 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Plain-torch mirrors of what both draft-step ops write besides their main output. + +The greedy draft token over the vocabulary-sharded logits, and the next draft step's state: position ids, +sequence and KV lengths, and the DSA decode metadata that the fused DSA decode-metadata kernel writes for a +generation-only batch. Shared by the references of ``glm_draft_sample_glue`` and ``glm_draft_ehproj_table``. +""" + +from __future__ import annotations + +from typing import Dict, List, Optional, Sequence + +import torch + + +def global_argmax(per_rank_logits: Sequence[torch.Tensor]) -> torch.Tensor: + """The argmax of the rank-ordered concatenation of the logits shards, in torch.argmax's order. + + A NaN beats every number and the first NaN wins; otherwise the largest value, the lowest index on ties. + Returns int32 ``[rows]``. + """ + logits = torch.cat([shard.float() for shard in per_rank_logits], dim=-1) + return torch.argmax(logits, dim=-1).to(torch.int32) + + +def _rms_normalize(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """``bf16((x * rsqrt(mean(x^2) + eps)) * weight)`` in FP32, one row per leading index.""" + rstd = torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + eps) + return ((x.float() * rstd) * weight.float()).to(torch.bfloat16) + + +def eh_proj_input_window( + tokens: torch.Tensor, + embed_table: torch.Tensor, + hidden: torch.Tensor, + enorm_weight: torch.Tensor, + hnorm_weight: torch.Tensor, + eps: float, + col_lo: int, + col_hi: int, +) -> torch.Tensor: + """Columns ``[col_lo, col_hi)`` of ``cat([enorm(embed(tokens)), hnorm(hidden)], -1)``. + + An id outside the table gives a zero embedding row. + """ + vocab = embed_table.shape[0] + ids = tokens.long() + valid = (ids >= 0) & (ids < vocab) + embeds = embed_table[ids.clamp(0, vocab - 1)].float() * valid.unsqueeze(1) + concat = torch.cat( + [_rms_normalize(embeds, enorm_weight, eps), _rms_normalize(hidden, hnorm_weight, eps)], + dim=-1, + ) + return concat[:, col_lo:col_hi].contiguous() + + +def next_step_state( + rows: int, + positions_in: Optional[torch.Tensor], + positions_out: Optional[torch.Tensor], + gather_ids: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + fill_seq_lens: bool, + kv_lens: Optional[torch.Tensor], + kv_update: int, + rewind_accepted: Optional[torch.Tensor], + draft_len: int, + block_offsets: Optional[torch.Tensor], + dsa_layout: List[int], + req_idx_per_token: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + gen_kv_indptr: Optional[torch.Tensor], + gen_cached_token_indptr: Optional[torch.Tensor], +) -> Dict[str, torch.Tensor]: + """What a draft-step op leaves in each state buffer it is handed, read as the buffers were before the call. + + Position ids: ``positions_out[r] = positions_in[gather_ids[r] if gather_ids else r] + 1``. Sequence lengths: + filled with 1 when ``fill_seq_lens``. KV lengths: ``+ 1`` (kv_update 1) or ``- (draft_len - + rewind_accepted[r])`` (kv_update 2). DSA decode metadata, from those lengths: the generation KV and cached + indptrs and, for each request's tokens (at most ``dsa_layout[1]`` per request, below ``dsa_layout[0]`` in + all), the request index and the indexer-K cache slot offsets, the block index truncated and clamped as the + fused DSA decode-metadata kernel computes it. Entries the op does not write keep their values. Returns the + expected buffers by argument name. + """ + out: Dict[str, torch.Tensor] = {} + if positions_out is not None: + expected = positions_out.clone() + source = gather_ids[:rows].long() if gather_ids is not None else torch.arange(rows) + expected[:rows] = positions_in.reshape(-1)[source.to(positions_in.device)] + 1 + out["positions_out"] = expected + seq = None + if seq_lens is not None: + expected = seq_lens.clone() + if fill_seq_lens: + expected[:rows] = 1 + out["seq_lens"] = expected + seq = [int(v) for v in expected[:rows].tolist()] + kv = None + if kv_lens is not None: + expected = kv_lens.clone() + if kv_update == 1: + expected[:rows] += 1 + elif kv_update == 2: + expected[:rows] -= draft_len - rewind_accepted[:rows].to(expected.dtype) + out["kv_lens"] = expected + kv = [int(v) for v in expected[:rows].tolist()] + if req_idx_per_token is None: + return out + + num_tokens, lanes, tokens_per_block, data_bytes, scale_size, block_stride, scale_base = ( + dsa_layout + ) + req_idx = req_idx_per_token.clone() + fp8 = slot_mapping_fp8.clone() + scale = slot_mapping_scale.clone() + kv_indptr = gen_kv_indptr.clone() + cached_indptr = gen_cached_token_indptr.clone() + kv_indptr[0] = 0 + cached_indptr[0] = 0 + blocks = block_offsets.cpu().tolist() + max_blocks = block_offsets.shape[1] + flat_start = 0 + kv_sum = 0 + cached_sum = 0 + for r in range(rows): + kv_sum += kv[r] + cached_sum += kv[r] - seq[r] + kv_indptr[1 + r] = kv_sum + cached_indptr[1 + r] = cached_sum + start = kv[r] - seq[r] + for j in range(min(seq[r], lanes)): + g = flat_start + j + if g >= num_tokens: + break + position = start + j + # Truncating division, a non-negative remainder, then the clamp. + block = abs(position) // tokens_per_block * (1 if position >= 0 else -1) + offset = position - block * tokens_per_block + if offset < 0: + offset += tokens_per_block + block = min(max(block, 0), max_blocks - 1) + block_id = blocks[r][block] + req_idx[g] = r + fp8[g] = block_id * block_stride + offset * data_bytes + scale[g] = block_id * block_stride + scale_base + offset * scale_size + flat_start += seq[r] + out.update( + req_idx_per_token=req_idx, + slot_mapping_fp8=fp8, + slot_mapping_scale=scale, + gen_kv_indptr=kv_indptr, + gen_cached_token_indptr=cached_indptr, + ) + return out + + +def assert_state_args_valid( + rows: int, + draft_tokens: torch.Tensor, + positions_in: Optional[torch.Tensor], + positions_out: Optional[torch.Tensor], + gather_ids: Optional[torch.Tensor], + kv_update: int, + draft_len: int, + dsa_layout: List[int], + req_idx_per_token: Optional[torch.Tensor], +) -> None: + """The state arguments both draft-step ops accept and answer wrongly.""" + if positions_out is not None and positions_in is not None: + # Lanes read positions_in[gather_ids[r]] while others write positions_out[r]: overlapping buffers race. + lo_in, hi_in = positions_in.data_ptr(), positions_in.data_ptr() + positions_in.nbytes + lo_out, hi_out = positions_out.data_ptr(), positions_out.data_ptr() + positions_out.nbytes + assert hi_in <= lo_out or hi_out <= lo_in, "positions_out must not overlap positions_in" + if gather_ids is not None: + source = gather_ids[:rows] + assert bool(((source >= 0) & (source < positions_in.numel())).all()), ( + "gather_ids must index positions_in" + ) + # The rewind subtracts draft_len - accepted from each KV length: another draft length rewinds wrongly. + assert kv_update != 2 or draft_len == draft_tokens.shape[1], ( + "the KV rewind needs draft_len equal to the step's draft length" + ) + if req_idx_per_token is not None: + # No token lanes, or no tokens per block, writes no metadata or divides by zero. + assert len(dsa_layout) == 7 and dsa_layout[1] >= 1 and dsa_layout[2] >= 1, ( + "dsa_layout needs at least one token lane and one token per block" + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_token_table.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_token_table.py new file mode 100644 index 000000000000..d36ab203accd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_token_table.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The constant operands of ``glm_draft_ehproj_table``, built once from the checkpoint. + +eh_proj is row-parallel over ``cat([enorm(embed(t)), hnorm(h)], -1)``, so its output splits into +``W[:, :H] @ enorm(embed(t))``, which depends on the draft token only, and ``W[:, H:] @ hnorm(h)``. Each rank +tabulates the first half for the tokens of its LM-head vocabulary shard and keeps its output rows of the second. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul + +from ._draft_glue_rmsnorm import draft_glue_rmsnorm_cat + +#: Table rows normalized and multiplied per block, and output columns per block (bounds the FP32 copies). +_ROW_BLOCK = 2048 +_COL_BLOCK = 1536 + + +def build_draft_token_table( + eh_proj_weight: torch.Tensor, + embed_table: torch.Tensor, + enorm_weight: torch.Tensor, + eps: float, + tp_rank: int, + tp_size: int, + vocab_shard: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """``(token_table, hidden_rows)`` of rank ``tp_rank``. + + ``eh_proj_weight`` is the whole ``[H, 2H]`` BF16 eh_proj weight and ``embed_table`` the whole ``[V, H]`` BF16 + embedding. ``token_table[t] = bf16(W[:, :H] @ enorm(embed_table[tp_rank * vocab_shard + t]))``: enorm by the + draft-step glue kernel (the arithmetic of draft step 0's eh_proj input), the product in FP32 without TF32, + one rounding; a token past the table gives a zero row. ``hidden_rows`` is this rank's ``[H / tp_size, H]`` + slice ``W[tp_rank * H / tp : (tp_rank + 1) * H / tp, H:]``, the weight of the hidden-half GEMM. + """ + hidden = embed_table.shape[1] + share = hidden // tp_size + device = embed_table.device + hidden_rows = eh_proj_weight[tp_rank * share : (tp_rank + 1) * share, hidden:].contiguous() + token_table = torch.empty((vocab_shard, hidden), dtype=torch.bfloat16, device=device) + first = tp_rank * vocab_shard + ids = torch.arange(first, first + vocab_shard, dtype=torch.int64, device=device) + # The glue kernel's hidden-row operand is not read for a window inside the embedding half. + unused_rows = torch.empty((_ROW_BLOCK, hidden), dtype=torch.bfloat16, device=device) + with true_fp32_matmul(): + for row in range(0, vocab_shard, _ROW_BLOCK): + n = min(_ROW_BLOCK, vocab_shard - row) + normed = draft_glue_rmsnorm_cat( + embed_table, + ids[row : row + n], + unused_rows[:n], + enorm_weight, + enorm_weight, + eps, + 0, + hidden, + False, + ).float() + for lo in range(0, hidden, _COL_BLOCK): + weight_t = eh_proj_weight[lo : lo + _COL_BLOCK, :hidden].float().t() + token_table[row : row + n, lo : lo + _COL_BLOCK] = (normed @ weight_t).to( + torch.bfloat16 + ) + return token_table, hidden_rows diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_gather_last_accepted_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_gather_last_accepted_rows.py new file mode 100644 index 000000000000..492124cbbf1d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_gather_last_accepted_rows.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Triton kernel behind ``gather_last_accepted_rows``: the step-0 row selection of the MTP draft loop. + +Request ``r`` selects row ``g[r] = batch_indices[r] * rows_per_request + num_accepted[r] - 1 + row_offset``, the +row of its last accepted token. One launch writes the int64 ``g`` and copies row ``g[r]`` of one or two row +tensors into row ``r`` of fresh outputs. +""" + +from __future__ import annotations + +from typing import Optional, Tuple, Union + +import torch +import triton +import triton.language as tl + +#: Columns one program copies. +_BLOCK = 1024 + + +@triton.jit +def _gather_last_accepted_rows_kernel( + src_ptr, + out_ptr, + src2_ptr, + out2_ptr, + batch_indices_ptr, + num_accepted_ptr, + gather_ids_ptr, + stride_src, + stride_out, + stride_src2, + stride_out2, + rows_per_request, + row_offset, + n_cols, + HAS_SECOND: tl.constexpr, + BLOCK: tl.constexpr, +): + req = tl.program_id(0) + chunk = tl.program_id(1) + row = ( + tl.load(batch_indices_ptr + req).to(tl.int64) * rows_per_request + + tl.load(num_accepted_ptr + req).to(tl.int64) + - 1 + + row_offset + ) + if chunk == 0: + tl.store(gather_ids_ptr + req, row) + cols = chunk * BLOCK + tl.arange(0, BLOCK) + mask = cols < n_cols + values = tl.load(src_ptr + row * stride_src + cols, mask=mask) + tl.store(out_ptr + req.to(tl.int64) * stride_out + cols, values, mask=mask) + if HAS_SECOND: + values2 = tl.load(src2_ptr + row * stride_src2 + cols, mask=mask) + tl.store(out2_ptr + req.to(tl.int64) * stride_out2 + cols, values2, mask=mask) + + +def gather_last_accepted_rows( + rows: torch.Tensor, + batch_indices: torch.Tensor, + num_accepted: torch.Tensor, + rows_per_request: int, + row_offset: int, + second: Optional[torch.Tensor] = None, +) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + """``(g, rows[g])``, plus ``second[g]`` as a third element when ``second`` is given, in one launch. + + ``rows`` and ``second`` are 2D with unit column stride; ``batch_indices`` and ``num_accepted`` are contiguous + integer vectors of one length. + """ + num_requests = batch_indices.shape[0] + n_cols = rows.shape[1] + gather_ids = torch.empty((num_requests,), dtype=torch.int64, device=rows.device) + out = rows.new_empty((num_requests, n_cols)) + out2 = second.new_empty((num_requests, second.shape[1])) if second is not None else out + if num_requests > 0: + _gather_last_accepted_rows_kernel[(num_requests, triton.cdiv(n_cols, _BLOCK))]( + rows, + out, + second if second is not None else rows, + out2, + batch_indices, + num_accepted, + gather_ids, + rows.stride(0), + out.stride(0), + second.stride(0) if second is not None else 0, + out2.stride(0), + rows_per_request, + row_offset, + n_cols, + HAS_SECOND=second is not None, + BLOCK=_BLOCK, + ) + if second is not None: + return gather_ids, out, out2 + return gather_ids, out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_ehproj_table.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_ehproj_table.py new file mode 100644 index 000000000000..2e7ae1fa500c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_ehproj_table.py @@ -0,0 +1,285 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""An MTP draft step that also produces the next step's eh_proj output, via glm_draft_ehproj_table.""" + +from __future__ import annotations + +from typing import Dict, List, Optional, Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._draft_step_state import assert_state_args_valid, global_argmax, next_step_state + + +def _step(**overrides) -> dict: + spec = dict( + rows=1, + draft_len=7, + step=1, + vocab_shard=9680, + logits="random", + positions="plain", + fill=False, + kv_update=1, + dsa=True, + ) + spec.update(overrides) + return spec + + +class GlmDraftEhprojTable(OpWrapper): + """The greedy draft token of every row and the next step's eh_proj output with its all-reduce and norm. + + eh_proj is row-parallel over `cat([enorm(embed(t)), hnorm(h)], -1)`, so its output splits into + `W[:, :H] @ enorm(embed(t))`, which depends on the token only, and `W[:, H:] @ hnorm(h)`, which is known before + the draft LM head. `token_table` holds the first half for the tokens of this rank's LM-head shard + (`[vocab_shard, H]` BF16, one rounding of the FP32 product); `hidden_half` is this rank's FP32 share + `[rows, H / group_size]` of the second half (columns `[rank * H / tp, (rank + 1) * H / tp)`). + + Collective over the TP group, in one 8-CTA cluster per rank: each rank finds its logits shard's argmax and + publishes, through the MNNVL workspace `exchange_buffer` / `exchange_flags`, the (index, value) packet, the + token-table row of its candidate and its hidden-half share. Every rank then takes the global argmax in + torch.argmax's order into column `draft_step` of `draft_tokens`, writes + `residual_out = x = bf16(token_table_of_winner[token] + hidden_half)` (one FP32 add) and + `normed_out = bf16((x * rsqrt(mean(x^2) + eps)) * norm_weight)`: what eh_proj with its all-reduce, + input_layernorm and a zero residual returns. The next step's position ids, lengths and DSA decode metadata are + written as `glm_draft_sample_glue` writes them. + + Under PDL the kernel reads only the norm weight before its grid-dependency wait and triggers its dependents + right after it; the hidden-half GEMM on the other stream must be joined by an event before the launch. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # spec: rows (requests: 1 at c1, at most 2 where the target builds the table), draft_len (K), step, + # vocab_shard (154880 / 16 at tp16), logits ("random" or "tied"), positions ("gather" on step 0, else + # "plain"), fill, kv_update, dsa. At c1 the op serves draft steps 0-5: step 0 turns the 8 verified rows into + # one (gathered positions, seq_lens filled with 1, KV rewound by the rejected drafts), steps 1-5 advance by one. + CELLS: tuple[Cell, ...] = ( + Cell( + why="c1 draft step 0: gathered positions, seq_lens fill, KV rewind, DSA metadata", + spec=_step(step=0, positions="gather", fill=True, kv_update=2), + ), + Cell( + why="c1 draft steps 1-5: positions + 1, KV + 1, DSA metadata", + spec=_step(step=3), + ), + Cell( + why="c1 middle step, every rank's maximum equal: rank 0's row wins", + spec=_step(step=2, logits="tied"), + ), + Cell( + why="two requests, the most rows the target builds the table for", + spec=_step(rows=2, step=1), + ), + ) + + note = """ + Serves 2 to 16 ranks whose shares of the 6144 hidden columns are whole 8-column vectors; at tp16 a CTA's 768 + columns span two ranks' 384-column shares, a mapping a 4-rank run does not reach. The workspace comes from + `DraftExchangeWorkspace` (one per op, built outside capture). + + `normed_out` differs from eh_proj + all-reduce + RMSNorm in the order of the norm's sum of squares (CTA + slices in order); `residual_out` differs from eh_proj's tp partial sums rounded to BF16 before their reduction: + here only the token half is rounded before x. Both feed the draft only. + """ + + def __call__( + self, + logits: torch.Tensor, + exchange_buffer: torch.Tensor, + exchange_flags: torch.Tensor, + group_size: int, + draft_tokens: torch.Tensor, + draft_step: int, + hidden_half: torch.Tensor, + token_table: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + normed_out: torch.Tensor, + residual_out: torch.Tensor, + positions_in: Optional[torch.Tensor], + positions_out: Optional[torch.Tensor], + gather_ids: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + fill_seq_lens: bool, + kv_lens: Optional[torch.Tensor], + kv_update: int, + rewind_accepted: Optional[torch.Tensor], + draft_len: int, + block_offsets: Optional[torch.Tensor], + dsa_layout: List[int], + req_idx_per_token: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + gen_kv_indptr: Optional[torch.Tensor], + gen_cached_token_indptr: Optional[torch.Tensor], + ) -> None: + torch.ops.trtllm.glm_draft_ehproj_table( + logits, + exchange_buffer, + exchange_flags, + group_size, + draft_tokens, + draft_step, + hidden_half, + token_table, + norm_weight, + eps, + normed_out, + residual_out, + positions_in, + positions_out, + gather_ids, + seq_lens, + fill_seq_lens, + kv_lens, + kv_update, + rewind_accepted, + draft_len, + block_offsets, + dsa_layout, + req_idx_per_token, + slot_mapping_fp8, + slot_mapping_scale, + gen_kv_indptr, + gen_cached_token_indptr, + ) + + def reference( + self, + logits: Sequence[torch.Tensor], + exchange_buffer: torch.Tensor, + exchange_flags: torch.Tensor, + group_size: int, + draft_tokens: torch.Tensor, + draft_step: int, + hidden_half: Sequence[torch.Tensor], + token_table: Sequence[torch.Tensor], + norm_weight: torch.Tensor, + eps: float, + normed_out: torch.Tensor, + residual_out: torch.Tensor, + positions_in: Optional[torch.Tensor], + positions_out: Optional[torch.Tensor], + gather_ids: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + fill_seq_lens: bool, + kv_lens: Optional[torch.Tensor], + kv_update: int, + rewind_accepted: Optional[torch.Tensor], + draft_len: int, + block_offsets: Optional[torch.Tensor], + dsa_layout: List[int], + req_idx_per_token: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + gen_kv_indptr: Optional[torch.Tensor], + gen_cached_token_indptr: Optional[torch.Tensor], + ) -> Dict[str, torch.Tensor]: + """What the op leaves in each buffer it writes, by argument name. + + The documented exception to mirroring `__call__`: `logits`, `hidden_half` and `token_table` are the + group's, in rank order, because every rank's output depends on all of them. Every other argument is the + calling rank's, read as it was before the call. The exchange workspace is not consulted. + """ + rows = logits[0].shape[0] + vocab_shard = logits[0].shape[1] + tokens = global_argmax(logits) + expected_draft = draft_tokens.clone() + expected_draft[:rows, draft_step] = tokens + token_half = torch.stack( + [ + token_table[int(t) // vocab_shard][int(t) % vocab_shard].float() + for t in tokens.tolist() + ] + ) + x = (token_half + torch.cat([h.float() for h in hidden_half], dim=-1)).to(torch.bfloat16) + rstd = torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + eps) + normed = ((x.float() * rstd) * norm_weight.float()).to(torch.bfloat16) + out = {"draft_tokens": expected_draft, "residual_out": x, "normed_out": normed} + out.update( + next_step_state( + rows, + positions_in, + positions_out, + gather_ids, + seq_lens, + fill_seq_lens, + kv_lens, + kv_update, + rewind_accepted, + draft_len, + block_offsets, + dsa_layout, + req_idx_per_token, + slot_mapping_fp8, + slot_mapping_scale, + gen_kv_indptr, + gen_cached_token_indptr, + ) + ) + return out + + def is_valid( + self, + logits: torch.Tensor, + exchange_buffer: torch.Tensor, + exchange_flags: torch.Tensor, + group_size: int, + draft_tokens: torch.Tensor, + draft_step: int, + hidden_half: torch.Tensor, + token_table: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + normed_out: torch.Tensor, + residual_out: torch.Tensor, + positions_in: Optional[torch.Tensor], + positions_out: Optional[torch.Tensor], + gather_ids: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + fill_seq_lens: bool, + kv_lens: Optional[torch.Tensor], + kv_update: int, + rewind_accepted: Optional[torch.Tensor], + draft_len: int, + block_offsets: Optional[torch.Tensor], + dsa_layout: List[int], + req_idx_per_token: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + gen_kv_indptr: Optional[torch.Tensor], + gen_cached_token_indptr: Optional[torch.Tensor], + ) -> None: + assert_state_args_valid( + logits.shape[0], + draft_tokens, + positions_in, + positions_out, + gather_ids, + kv_update, + draft_len, + dsa_layout, + req_idx_per_token, + ) + + # x is stored row by row while the cluster still reads the token table for the later rows, and the norm + # after the last one: an output overlapping an input, or the other output, replaces values the kernel has + # yet to read or to return. + def overlaps(a: torch.Tensor, b: torch.Tensor) -> bool: + return a.data_ptr() < b.data_ptr() + b.nbytes and b.data_ptr() < a.data_ptr() + a.nbytes + + for name, t in (("normed_out", normed_out), ("residual_out", residual_out)): + assert not any(overlaps(t, other) for other in (logits, hidden_half, token_table)), ( + f"{name} must not overlap the op's inputs" + ) + assert not overlaps(normed_out, residual_out), ( + "normed_out and residual_out must be separate buffers" + ) + + +glm_draft_ehproj_table = GlmDraftEhprojTable() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_glue_rmsnorm_cat.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_glue_rmsnorm_cat.py new file mode 100644 index 000000000000..45dfe48fb242 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_glue_rmsnorm_cat.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""This rank's MTP eh_proj input window -- embedding gather, enorm, hnorm, concat, column split -- in one kernel.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._draft_glue_rmsnorm import draft_glue_rmsnorm_cat as _draft_glue_rmsnorm_cat + +#: GLM-5.2's hidden size and the eh_proj input width (the concat of two hidden rows). +_HIDDEN = 6144 +#: Columns of the eh_proj input each tp16 rank holds (row-parallel eh_proj). +_TP16_WINDOW = 2 * _HIDDEN // 16 + + +def _tp16_window(rank: int) -> tuple[int, int]: + return rank * _TP16_WINDOW, (rank + 1) * _TP16_WINDOW + + +class GlmDraftGlueRmsnormCat(OpWrapper): + """Columns `[col_lo, col_hi)` of `cat([rmsnorm(embed_table[input_ids]), rmsnorm(hidden_states)], -1)`. + + The embedding rows are normalized with `enorm_weight`, the hidden rows with `hnorm_weight`, both with `eps`: + `bf16((x * rsqrt(mean(x^2) + eps)) * w)` in FP32. An id outside `[0, embed_table.shape[0])` gives a zero + embedding row, as the Embedding module does. Only the halves the window overlaps are computed and only the + window is stored. Returns a contiguous BF16 `[M, col_hi - col_lo]` tensor. + + With `enable_pdl` the kernel waits for its producer before any load and lets its dependent (eh_proj) launch + right after that wait. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # Draft step 0 at c1: the 8 verified rows of the one request (accepted-token ids, the target's last hidden + # rows). eh_proj is row-parallel over the 12288-column concat, so rank r feeds it columns + # [768 r, 768 (r + 1)): ranks 0-7 normalize only the embedding half, ranks 8-15 only the hidden half. The + # spec's `vocab` is the table's row count; the gather does not depend on it. + CELLS: tuple[Cell, ...] = ( + Cell( + why="c1 step 0, tp16 rank 0: the first embedding-half window", + spec=dict(rows=8, vocab=4096, window=_tp16_window(0), ids="valid"), + ), + Cell( + why="c1 step 0, tp16 rank 7: the last embedding-half window", + spec=dict(rows=8, vocab=4096, window=_tp16_window(7), ids="valid"), + ), + Cell( + why="c1 step 0, tp16 rank 8: the first hidden-half window", + spec=dict(rows=8, vocab=4096, window=_tp16_window(8), ids="valid"), + ), + Cell( + why="c1 step 0, tp16 rank 15: the last hidden-half window", + spec=dict(rows=8, vocab=4096, window=_tp16_window(15), ids="valid"), + ), + Cell( + why="ids outside the table (-1 padding, past the end) give zero rows", + spec=dict(rows=8, vocab=4096, window=_tp16_window(3), ids="outside"), + ), + Cell( + why="one row: a draft step that feeds the glue a single token", + spec=dict(rows=1, vocab=4096, window=_tp16_window(5), ids="valid"), + ), + ) + + note = """ + The kernel covers hidden sizes whose FlashInfer RMSNormKernel runs one CTA per row (6144 here); a window over + both halves (tp below 2, or attention DP) is supported but no target of this family reaches it. Each row is + normalized with FlashInfer's CuTe RMSNorm thread geometry, reduction and arithmetic statements; the cells + certify it against the FP32 reference. + """ + + def __call__( + self, + embed_table: torch.Tensor, + input_ids: torch.Tensor, + hidden_states: torch.Tensor, + enorm_weight: torch.Tensor, + hnorm_weight: torch.Tensor, + eps: float, + col_lo: int, + col_hi: int, + enable_pdl: bool, + ) -> torch.Tensor: + return _draft_glue_rmsnorm_cat( + embed_table, + input_ids, + hidden_states, + enorm_weight, + hnorm_weight, + eps, + col_lo, + col_hi, + enable_pdl, + ) + + def reference( + self, + embed_table: torch.Tensor, + input_ids: torch.Tensor, + hidden_states: torch.Tensor, + enorm_weight: torch.Tensor, + hnorm_weight: torch.Tensor, + eps: float, + col_lo: int, + col_hi: int, + enable_pdl: bool, + ) -> torch.Tensor: + """The unfused chain in FP32: gather with the zero-row guard, both RMSNorms, concat, column slice.""" + vocab = embed_table.shape[0] + ids = input_ids.long() + valid = (ids >= 0) & (ids < vocab) + embeds = embed_table[ids.clamp(0, vocab - 1)].float() * valid.unsqueeze(1) + + def rmsnorm(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) + return ((x * rstd) * weight.float()).to(torch.bfloat16) + + concat = torch.cat( + [rmsnorm(embeds, enorm_weight), rmsnorm(hidden_states.float(), hnorm_weight)], dim=-1 + ) + return concat[:, col_lo:col_hi].contiguous() + + def is_valid( + self, + embed_table: torch.Tensor, + input_ids: torch.Tensor, + hidden_states: torch.Tensor, + enorm_weight: torch.Tensor, + hnorm_weight: torch.Tensor, + eps: float, + col_lo: int, + col_hi: int, + enable_pdl: bool, + ) -> None: + hidden = hidden_states.shape[1] + # The predicate stores whole 8-column vectors by their first column: a window edge inside a vector moves + # columns into or out of the window. + assert 0 <= col_lo < col_hi <= 2 * hidden and col_lo % 8 == 0 and col_hi % 8 == 0, ( + "the window must be 8-column aligned inside [0, 2 * hidden)" + ) + # The kernel is compiled for 16-byte aligned rows whose stride is a multiple of 8 elements. + rows = [("hidden_states", hidden_states)] + if col_lo < hidden: + rows.append(("embed_table", embed_table)) + assert input_ids.dim() == 1 and input_ids.is_contiguous(), ( + "input_ids must be a contiguous vector" + ) + assert input_ids.shape[0] == hidden_states.shape[0], "one id per hidden row" + for name, t in rows: + assert t.dim() == 2 and t.shape[1] == hidden and t.stride(1) == 1, ( + f"{name} must be [rows, hidden] with unit column stride" + ) + assert t.stride(0) % 8 == 0 and t.data_ptr() % 16 == 0, ( + f"{name} rows must be 16-byte aligned" + ) + for name, w in (("enorm_weight", enorm_weight), ("hnorm_weight", hnorm_weight)): + assert w.shape == (hidden,) and w.is_contiguous() and w.data_ptr() % 16 == 0, ( + f"{name} must be a contiguous, 16-byte aligned [hidden] vector" + ) + + +glm_draft_glue_rmsnorm_cat = GlmDraftGlueRmsnormCat() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_sample_glue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_sample_glue.py new file mode 100644 index 000000000000..741c387dedcf --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_sample_glue.py @@ -0,0 +1,322 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One MTP draft step after the vocabulary-sharded draft LM head, in one kernel, via glm_draft_sample_glue.""" + +from __future__ import annotations + +from typing import Dict, List, Optional, Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._draft_step_state import ( + assert_state_args_valid, + eh_proj_input_window, + global_argmax, + next_step_state, +) + + +def _step(**overrides) -> dict: + spec = dict( + rows=1, + draft_len=7, + step=6, + vocab_shard=9680, + logits="random", + glue=False, + positions=None, + fill=False, + kv_update=1, + dsa=True, + next_new=False, + ) + spec.update(overrides) + return spec + + +class GlmDraftSampleGlue(OpWrapper): + """The greedy draft token of every row, and optionally the next step's inputs and state, in one CTA. + + Collective over the TP group (`group_size` ranks, the workspace's group): each rank takes the argmax of its + logits shard `[rows, vocab_shard]` (rank r holds tokens `[r * vocab_shard, (r + 1) * vocab_shard)`), the ranks + exchange the (index, value) pairs through the MNNVL workspace `exchange_buffer` / `exchange_flags`, and every + rank writes the global argmax, in torch.argmax's order, into column `draft_step` of `draft_tokens`. Then, each + part skipped when its output is None: + - `glue_out`: this rank's window `[col_lo, col_hi)` of `cat([enorm(embed_table[token]), hnorm(hidden)], -1)`, + the next step's eh_proj input (FlashInfer's CuTe RMSNorm arithmetic), contiguous; + - `positions_out[r] = positions_in[gather_ids[r] if gather_ids is not None else r] + 1`; + - `seq_lens` filled with 1 when `fill_seq_lens`; `kv_lens` advanced as `kv_update` says (0 none, 1 `+ 1`, + 2 `- (draft_len - rewind_accepted)`); + - the DSA decode metadata of the next step from those lengths (`req_idx_per_token`, the two slot mappings, the + two generation indptrs; `dsa_layout` = [tokens, token lanes per request, tokens per block, data bytes per + token, scale bytes per token, bytes per block, scale offset in a block]); + - `next_new_tokens[r] = [accepted_tokens[batch_indices[r], num_accepted[r] - 1], draft tokens...]`, the + output rows after the last draft step. + + Under PDL the kernel reads only the norm weights before its grid-dependency wait and triggers its dependents + right after it. The workspace serves this op only (each call empties the buffer of the previous one). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # spec: rows (requests), draft_len (K), step, vocab_shard (the per-rank LM-head shard: 154880 / 16 at tp16), + # logits ("random", or "tied": every rank's maximum equal), glue (the eh_proj input window, this rank's + # 2 * 6144 / tp columns), positions (None, "gather" or "plain"), fill (seq_lens), kv_update, dsa (metadata), + # next_new (output rows). At c1 the token table serves draft steps 0-5 and this op the last one (step 6): + # token, KV + 1, DSA metadata and the next-new-tokens rows. Above two requests per rank (no token table) it + # serves every step, with the eh_proj input window. + CELLS: tuple[Cell, ...] = ( + Cell( + why="c1 last draft step: token, KV + 1, DSA metadata, next-new-tokens rows", + spec=_step(next_new=True), + ), + Cell( + why="c1 last draft step, every rank's maximum equal: rank 0's index wins", + spec=_step(logits="tied", next_new=True), + ), + Cell( + why="step 0 without the token table: window, gathered positions, fill, rewind", + spec=_step(rows=4, step=0, glue=True, positions="gather", fill=True, kv_update=2), + ), + Cell( + why="a middle step without the token table: window, positions + 1, KV + 1", + spec=_step(rows=4, step=3, glue=True, positions="plain"), + ), + Cell( + why="the most rows one launch serves (8 requests), last step", + spec=_step(rows=8, next_new=True), + ), + ) + + note = """ + The workspace comes from `DraftExchangeWorkspace` (one per op, built outside capture). Ranks must call in + lockstep with the same `rows`; a missing rank leaves the others polling. + + The vocabulary shards must be equal and in rank order, which an unpadded vocabulary-parallel LM head gives + (154880 = 16 x 9680). + """ + + def __call__( + self, + logits: torch.Tensor, + exchange_buffer: torch.Tensor, + exchange_flags: torch.Tensor, + group_size: int, + draft_tokens: torch.Tensor, + draft_step: int, + embed_table: Optional[torch.Tensor], + hidden: Optional[torch.Tensor], + enorm_weight: Optional[torch.Tensor], + hnorm_weight: Optional[torch.Tensor], + eps: float, + col_lo: int, + col_hi: int, + glue_out: Optional[torch.Tensor], + positions_in: Optional[torch.Tensor], + positions_out: Optional[torch.Tensor], + gather_ids: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + fill_seq_lens: bool, + kv_lens: Optional[torch.Tensor], + kv_update: int, + rewind_accepted: Optional[torch.Tensor], + draft_len: int, + block_offsets: Optional[torch.Tensor], + dsa_layout: List[int], + req_idx_per_token: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + gen_kv_indptr: Optional[torch.Tensor], + gen_cached_token_indptr: Optional[torch.Tensor], + accepted_tokens: Optional[torch.Tensor], + batch_indices: Optional[torch.Tensor], + num_accepted: Optional[torch.Tensor], + next_new_tokens: Optional[torch.Tensor], + ) -> None: + torch.ops.trtllm.glm_draft_sample_glue( + logits, + exchange_buffer, + exchange_flags, + group_size, + draft_tokens, + draft_step, + embed_table, + hidden, + enorm_weight, + hnorm_weight, + eps, + col_lo, + col_hi, + glue_out, + positions_in, + positions_out, + gather_ids, + seq_lens, + fill_seq_lens, + kv_lens, + kv_update, + rewind_accepted, + draft_len, + block_offsets, + dsa_layout, + req_idx_per_token, + slot_mapping_fp8, + slot_mapping_scale, + gen_kv_indptr, + gen_cached_token_indptr, + accepted_tokens, + batch_indices, + num_accepted, + next_new_tokens, + ) + + def reference( + self, + logits: Sequence[torch.Tensor], + exchange_buffer: torch.Tensor, + exchange_flags: torch.Tensor, + group_size: int, + draft_tokens: torch.Tensor, + draft_step: int, + embed_table: Optional[torch.Tensor], + hidden: Optional[torch.Tensor], + enorm_weight: Optional[torch.Tensor], + hnorm_weight: Optional[torch.Tensor], + eps: float, + col_lo: int, + col_hi: int, + glue_out: Optional[torch.Tensor], + positions_in: Optional[torch.Tensor], + positions_out: Optional[torch.Tensor], + gather_ids: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + fill_seq_lens: bool, + kv_lens: Optional[torch.Tensor], + kv_update: int, + rewind_accepted: Optional[torch.Tensor], + draft_len: int, + block_offsets: Optional[torch.Tensor], + dsa_layout: List[int], + req_idx_per_token: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + gen_kv_indptr: Optional[torch.Tensor], + gen_cached_token_indptr: Optional[torch.Tensor], + accepted_tokens: Optional[torch.Tensor], + batch_indices: Optional[torch.Tensor], + num_accepted: Optional[torch.Tensor], + next_new_tokens: Optional[torch.Tensor], + ) -> Dict[str, torch.Tensor]: + """What the op leaves in each buffer it writes, by argument name. + + The documented exception to mirroring `__call__`: the token is a function of every rank's logits, so + `logits` is the group's shards in rank order. Every other argument is the calling rank's, and the buffers + the op updates in place are read as they were before the call. The exchange workspace is not consulted. + """ + rows = logits[0].shape[0] + tokens = global_argmax(logits) + expected_draft = draft_tokens.clone() + expected_draft[:rows, draft_step] = tokens + out = {"draft_tokens": expected_draft} + if glue_out is not None: + out["glue_out"] = eh_proj_input_window( + tokens, embed_table, hidden, enorm_weight, hnorm_weight, eps, col_lo, col_hi + ) + out.update( + next_step_state( + rows, + positions_in, + positions_out, + gather_ids, + seq_lens, + fill_seq_lens, + kv_lens, + kv_update, + rewind_accepted, + draft_len, + block_offsets, + dsa_layout, + req_idx_per_token, + slot_mapping_fp8, + slot_mapping_scale, + gen_kv_indptr, + gen_cached_token_indptr, + ) + ) + if next_new_tokens is not None: + expected = next_new_tokens.clone().view(rows, -1) + last = accepted_tokens[batch_indices[:rows].long(), num_accepted[:rows].long() - 1] + expected[:, 0] = last + expected[:, 1:] = expected_draft[:rows] + out["next_new_tokens"] = expected.view_as(next_new_tokens) + return out + + def is_valid( + self, + logits: torch.Tensor, + exchange_buffer: torch.Tensor, + exchange_flags: torch.Tensor, + group_size: int, + draft_tokens: torch.Tensor, + draft_step: int, + embed_table: Optional[torch.Tensor], + hidden: Optional[torch.Tensor], + enorm_weight: Optional[torch.Tensor], + hnorm_weight: Optional[torch.Tensor], + eps: float, + col_lo: int, + col_hi: int, + glue_out: Optional[torch.Tensor], + positions_in: Optional[torch.Tensor], + positions_out: Optional[torch.Tensor], + gather_ids: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + fill_seq_lens: bool, + kv_lens: Optional[torch.Tensor], + kv_update: int, + rewind_accepted: Optional[torch.Tensor], + draft_len: int, + block_offsets: Optional[torch.Tensor], + dsa_layout: List[int], + req_idx_per_token: Optional[torch.Tensor], + slot_mapping_fp8: Optional[torch.Tensor], + slot_mapping_scale: Optional[torch.Tensor], + gen_kv_indptr: Optional[torch.Tensor], + gen_cached_token_indptr: Optional[torch.Tensor], + accepted_tokens: Optional[torch.Tensor], + batch_indices: Optional[torch.Tensor], + num_accepted: Optional[torch.Tensor], + next_new_tokens: Optional[torch.Tensor], + ) -> None: + rows = logits.shape[0] + assert_state_args_valid( + rows, + draft_tokens, + positions_in, + positions_out, + gather_ids, + kv_update, + draft_len, + dsa_layout, + req_idx_per_token, + ) + if next_new_tokens is not None: + # The kernel writes row r at offset r * (draft_len + 1) and reads + # accepted_tokens[batch_indices[r], num_accepted[r] - 1] unbounded. + assert tuple(next_new_tokens.shape) == (rows, draft_tokens.shape[1] + 1), ( + "next_new_tokens must be exactly [rows, draft_len + 1]" + ) + index = batch_indices[:rows].long() + count = num_accepted[:rows].long() + assert bool(((index >= 0) & (index < accepted_tokens.shape[0])).all()), ( + "batch_indices must index accepted_tokens" + ) + assert bool(((count >= 1) & (count <= accepted_tokens.shape[1])).all()), ( + "num_accepted must lie in [1, accepted_tokens.shape[1]]" + ) + + +glm_draft_sample_glue = GlmDraftSampleGlue() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/gather_last_accepted_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/gather_last_accepted_rows.py new file mode 100644 index 000000000000..d84da2985147 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/gather_last_accepted_rows.py @@ -0,0 +1,123 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MTP draft step 0's row selection -- each request's last accepted row, and its index -- in one Triton launch.""" + +from __future__ import annotations + +from typing import Optional, Tuple, Union + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._gather_last_accepted_rows import gather_last_accepted_rows as _gather_last_accepted_rows + + +class GlmGatherLastAcceptedRows(OpWrapper): + """Each request's last accepted row of `rows` (and of `second`), and the row's index. + + Request `r` selects `g[r] = batch_indices[r] * rows_per_request + num_accepted[r] - 1 + row_offset`. Draft + step 0 of the MTP loop runs the layer on every verified row of a request (`K + 1` of them), but only the row of + the request's last accepted token feeds the LM head and the next step. Returns `(g, rows[g])`, or + `(g, rows[g], second[g])`: `g` int64 `[R]`, the rows fresh `[R, n]` tensors. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # At c1 one request verifies K + 1 = 8 rows and step 0 selects one: after the post-attention all-reduce, the + # normalized hidden rows and the residual together (the MoE, its all-reduce and the LM head then run on that + # row), or the target's hidden rows alone where no row select applies. Generation-only batches only, so the + # row offset (the context tokens ahead of them) is 0. + CELLS: tuple[Cell, ...] = ( + Cell( + why="c1 step-0 row select: hidden and residual rows [8, 6144], one request", + spec=dict(requests=1, rows_per_request=8, cols=6144, second=True, accepted=(3,)), + ), + Cell( + why="c1 step-0 gather of the target hidden rows alone, every draft accepted", + spec=dict(requests=1, rows_per_request=8, cols=6144, second=False, accepted=(8,)), + ), + Cell( + why="four requests in a permuted batch order, first rows accepted", + spec=dict( + requests=4, + rows_per_request=8, + cols=6144, + second=True, + accepted=(1, 8, 4, 2), + batch_order=(2, 0, 3, 1), + ), + ), + ) + + note = """ + The kernel reads `batch_indices` and `num_accepted` as dense vectors and does not bound `g`: a count outside + [1, rows_per_request] reads another request's row or memory outside `rows`. + """ + + def __call__( + self, + rows: torch.Tensor, + batch_indices: torch.Tensor, + num_accepted: torch.Tensor, + rows_per_request: int, + row_offset: int, + second: Optional[torch.Tensor] = None, + ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + return _gather_last_accepted_rows( + rows, batch_indices, num_accepted, rows_per_request, row_offset, second + ) + + def reference( + self, + rows: torch.Tensor, + batch_indices: torch.Tensor, + num_accepted: torch.Tensor, + rows_per_request: int, + row_offset: int, + second: Optional[torch.Tensor] = None, + ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + """The draft loop's index chain and its row gathers.""" + gather_ids = batch_indices.long() * rows_per_request + num_accepted.long() - 1 + row_offset + if second is None: + return gather_ids, rows[gather_ids] + return gather_ids, rows[gather_ids], second[gather_ids] + + def is_valid( + self, + rows: torch.Tensor, + batch_indices: torch.Tensor, + num_accepted: torch.Tensor, + rows_per_request: int, + row_offset: int, + second: Optional[torch.Tensor] = None, + ) -> None: + # The kernel addresses rows through stride(0) with unit column stride, and the index vectors as dense + # arrays; any other layout is read from the wrong elements. + assert rows.dim() == 2 and rows.stride(1) == 1, "rows must be 2D with unit column stride" + assert second is None or (second.dim() == 2 and second.stride(1) == 1), ( + "second must be 2D with unit column stride" + ) + # The kernel copies rows.shape[1] columns of both tensors: a wider second keeps garbage in its tail. + assert second is None or second.shape[1] == rows.shape[1], ( + "second must have the columns of rows" + ) + assert batch_indices.dim() == 1 and batch_indices.is_contiguous(), ( + "batch_indices must be a contiguous vector" + ) + assert num_accepted.dim() == 1 and num_accepted.is_contiguous(), ( + "num_accepted must be a contiguous vector" + ) + assert num_accepted.shape[0] >= batch_indices.shape[0], "one accepted count per request" + # The selected row is not bounded on the device. + count = num_accepted[: batch_indices.shape[0]] + assert bool(((count >= 1) & (count <= rows_per_request)).all()), ( + "num_accepted must lie in [1, rows_per_request]" + ) + gather_ids = batch_indices.long() * rows_per_request + count.long() - 1 + row_offset + assert bool(((gather_ids >= 0) & (gather_ids < rows.shape[0])).all()), ( + "the selected rows must lie inside rows" + ) + + +glm_gather_last_accepted_rows = GlmGatherLastAcceptedRows() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/verify_head_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/verify_head_accept.py new file mode 100644 index 000000000000..e8acff1335e4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/verify_head_accept.py @@ -0,0 +1,238 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Strict draft acceptance with the forced accepted-token count in one kernel, via the glm_verify_head_accept op.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class GlmVerifyHeadAccept(OpWrapper): + """The accepted-token counts of one verify step, and optionally its accepted-token rows. + + A context request accepts its one sampled token (count 1). A generation request with sampled target tokens + `t[0..K]` and draft tokens `d[0..K-1]` accepts `1 + (length of the leading run of d[j] == t[j])`, then the + forced count replaces that as `force_mode` says: + - 0: no force; + - 1: `force_base`; + - 2: `min(force_base + (force_pool[i] < force_frac), force_max)` with + `i = (c * counter_stride + g * slot_stride) & (pool_size - 1)`, `g` the generation index and `c` the value + of `force_counter` after this call adds one to it (the comparison in FP32). + + That is `SpecWorkerBase._sample_and_accept_draft_tokens_base` after sampling with + `_apply_force_accepted_tokens`, for the padded layout (`target_tokens` holds one token per context request, + then `K + 1` per generation request). Returns the counts `[num_contexts + num_gens]` int32. With + `accepted_tokens` the op also writes the accepted-token rows: column 0 of each context row, the whole row of + each generation request (its target tokens); without it the caller reads the generation rows from + `target_tokens` itself. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # spec: num_contexts, num_gens, draft_len (K = 7 at c1), target dtype, matches (the leading draft tokens equal + # to the targets, per generation request), force (mode, base, max, frac), rows (accepted_tokens given). The + # hash constants are the spec worker's: counter stride 6007, slot stride 1009, pool 2^16. + # At c1 the gen server verifies one request with K = 7 and the cell config forces 2.91 accepted draft tokens: + # mode 2, base 3, max 8, frac 0.91; the accepted rows alias the sampled tokens. + CELLS: tuple[Cell, ...] = ( + Cell( + why="c1 verify step: one request, K 7, forced 2.91 (base 3, frac 0.91)", + spec=dict( + num_contexts=0, + num_gens=1, + draft_len=7, + target_dtype=torch.int32, + matches=(2,), + force=(2, 3, 8, 0.91), + rows=False, + ), + ), + Cell( + why="c1 verify step without a forced count: the strict count", + spec=dict( + num_contexts=0, + num_gens=1, + draft_len=7, + target_dtype=torch.int32, + matches=(5,), + force=(0, 0, 0, 0.0), + rows=False, + ), + ), + Cell( + why="every draft token accepted: the count reaches K + 1", + spec=dict( + num_contexts=0, + num_gens=1, + draft_len=7, + target_dtype=torch.int32, + matches=(7,), + force=(0, 0, 0, 0.0), + rows=True, + ), + ), + Cell( + why="an integral forced count (3.0): mode 1, every request gets base 4", + spec=dict( + num_contexts=0, + num_gens=1, + draft_len=7, + target_dtype=torch.int32, + matches=(0,), + force=(1, 4, 8, 0.0), + rows=False, + ), + ), + Cell( + why="warm-up batch with a context request, int64 targets, rows written", + spec=dict( + num_contexts=1, + num_gens=1, + draft_len=7, + target_dtype=torch.int64, + matches=(3,), + force=(2, 3, 8, 0.91), + rows=True, + ), + ), + Cell( + why="several generation requests: the pool index follows the slot", + spec=dict( + num_contexts=0, + num_gens=4, + draft_len=7, + target_dtype=torch.int32, + matches=(0, 7, 2, 4), + force=(2, 3, 8, 0.5), + rows=True, + ), + ), + ) + + note = """ + Not for ragged verification: a request that verified fewer than K + 1 positions has no padded row, and the op + would read its neighbour's tokens. The caller keeps the ATen chain there. + + The kernel waits for its producer (the sampler) before it reads anything and lets its dependents launch at its + end. Mode 2 advances `force_counter` once per call, also inside a captured graph's replays, exactly as the ATen + chain's in-place add does, so the op and the chain draw the same pool entries given the same counter. + """ + + def __call__( + self, + target_tokens: torch.Tensor, + draft_tokens: torch.Tensor, + num_contexts: int, + accepted_tokens: Optional[torch.Tensor], + force_pool: Optional[torch.Tensor], + force_counter: Optional[torch.Tensor], + force_mode: int, + force_base: int, + force_max: int, + force_frac: float, + counter_stride: int, + slot_stride: int, + pool_size: int, + ) -> torch.Tensor: + return torch.ops.trtllm.glm_verify_head_accept( + target_tokens, + draft_tokens, + num_contexts, + accepted_tokens, + force_pool, + force_counter, + force_mode, + force_base, + force_max, + force_frac, + counter_stride, + slot_stride, + pool_size, + ) + + def reference( + self, + target_tokens: torch.Tensor, + draft_tokens: torch.Tensor, + num_contexts: int, + accepted_tokens: Optional[torch.Tensor], + force_pool: Optional[torch.Tensor], + force_counter: Optional[torch.Tensor], + force_mode: int, + force_base: int, + force_max: int, + force_frac: float, + counter_stride: int, + slot_stride: int, + pool_size: int, + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + """The ATen chain of the one-model spec worker: `(counts, rows, counter)`. + + `rows` is what `accepted_tokens` holds after the call (its unwritten entries keep the values it was handed) + and `counter` what `force_counter` holds after it; each is None when the op is not handed that buffer + (or, for the counter, does not advance it). Both buffers are read as they were before the call. + """ + num_gens, draft_len = draft_tokens.shape + batch_size = num_contexts + num_gens + target = target_tokens.reshape(-1) + counts = torch.ones(batch_size, dtype=torch.int32, device=draft_tokens.device) + gen_target = target[num_contexts : num_contexts + num_gens * (draft_len + 1)].reshape( + num_gens, draft_len + 1 + ) + counts[num_contexts:] += ( + torch.cumprod((draft_tokens.long() == gen_target[:, :draft_len].long()).int(), dim=-1) + .sum(1) + .int() + ) + counter = None + if force_mode == 1: + counts[num_contexts:] = force_base + elif force_mode == 2: + counter = force_counter + 1 + slots = torch.arange(num_gens, dtype=torch.int64, device=draft_tokens.device) + index = (counter * counter_stride + slots * slot_stride) & (pool_size - 1) + extra = (force_pool[index] < force_frac).int() + counts[num_contexts:] = (force_base + extra).clamp(max=force_max).int() + rows = None + if accepted_tokens is not None: + rows = accepted_tokens.clone() + rows[:num_contexts, 0] = target[:num_contexts].int() + rows[num_contexts:batch_size] = gen_target.int() + return counts, rows, counter + + def is_valid( + self, + target_tokens: torch.Tensor, + draft_tokens: torch.Tensor, + num_contexts: int, + accepted_tokens: Optional[torch.Tensor], + force_pool: Optional[torch.Tensor], + force_counter: Optional[torch.Tensor], + force_mode: int, + force_base: int, + force_max: int, + force_frac: float, + counter_stride: int, + slot_stride: int, + pool_size: int, + ) -> None: + num_gens, draft_len = draft_tokens.shape + # The kernel writes row r at offset r * (K + 1): a wider buffer (rows of the maximum draft length while + # this step verifies fewer) is accepted by the op's size check and filled with the wrong row stride. + assert accepted_tokens is None or tuple(accepted_tokens.shape) == ( + num_contexts + num_gens, + draft_len + 1, + ), "accepted_tokens must be exactly [num_contexts + num_gens, draft_len + 1]" + # A forced count outside [1, K + 1] is written as given and sends the caller past its draft window. + if force_mode != 0: + assert 1 <= force_base <= force_max == draft_len + 1, ( + "a forced count needs 1 <= force_base <= force_max == draft_len + 1" + ) + + +glm_verify_head_accept = GlmVerifyHeadAccept() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py new file mode 100644 index 000000000000..8d34e6bf535e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GLM-5.2-NVFP4 checkpoint tensors mapped onto the tp16 target's parameters. + +The target's modules are the built-in GLM-5.2 modules (see ``modeling.py``), +so the checkpoint loads through the built-in loader into exactly the shapes +and layouts the context server uses. The MTP draft steps also derive operands +from two of the MTP layer's checkpoint tensors (``DraftWeights``), taken before +the built-in loader shards them. +""" + +from __future__ import annotations + +import copy +from dataclasses import dataclass +from typing import Optional + +import torch +from transformers import PretrainedConfig + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3WeightLoader +from tensorrt_llm._torch.modules.linear import TensorParallelMode, load_weight_shard + + +def normalize_model_config( + model_config: ModelConfig[PretrainedConfig], +) -> ModelConfig[PretrainedConfig]: + """The two config rewrites the built-in implementation applies before building modules. + + - Legacy checkpoints name the fused q_a / kv_a projection ``fused_a`` in their + per-module quant config; the modules call it ``kv_a_proj_with_mqa``. + - GLM-5.2 publishes ``model_type = glm_moe_dsa``; its attention is DeepSeek-V3.2's + (MLA with the DSA indexer), which the modules select by ``deepseek_v32``. + """ + if model_config.quant_config_dict is not None: + model_config = copy.deepcopy(model_config) + quant_config_dict = {} + for key, val in model_config.quant_config_dict.items(): + key_split = key.split(".") + if key_split[-1] == "fused_a": + key = ".".join(key_split[:-1] + ["kv_a_proj_with_mqa"]) + quant_config_dict[key] = val + model_config._frozen = False + model_config.quant_config_dict = quant_config_dict + model_config._frozen = True + if model_config.pretrained_config.model_type == "glm_moe_dsa": + model_config._frozen = False + model_config.pretrained_config.model_type = "deepseek_v32" + model_config._frozen = True + return model_config + + +@dataclass(frozen=True, eq=False) +class DraftWeights: + """The MTP layer's checkpoint tensors the decode target derives draft-step operands from. + + Attributes: + eh_proj: The whole ``[H, 2H]`` eh_proj weight, on the model's device. The built-in + eh_proj is row-parallel (``2H / tp`` input columns per rank); the draft token table + needs every column. + w3_w1: ``[E, 2 * I / tp, H]``: per expert, this rank's ``I / tp`` rows of up_proj, then + the same rows of gate_proj (the low-M BF16 MoE's layout). + w2: ``[E, H, I / tp]``: per expert, this rank's ``I / tp`` columns of down_proj. + + ``w3_w1`` and ``w2`` hold the slices the built-in MoE loader gives this rank, in plain + row-major layout: the built-in BF16 MoE may keep its own copy in a backend layout. + """ + + eh_proj: torch.Tensor + w3_w1: torch.Tensor + w2: torch.Tensor + + +def load_weights(model, weights, draft: bool) -> Optional[DraftWeights]: + """Load the checkpoint into the target's modules with the built-in loader. + + With ``draft``, first take the MTP layer's ``DraftWeights`` (the built-in loader releases + checkpoint tensors once it has copied them) and return them; None without ``draft`` or + when the checkpoint has no MTP layer. + """ + draft_weights = _draft_weights(model, weights) if draft else None + DeepseekV3WeightLoader(model).load_weights(weights) + return draft_weights + + +def _draft_weights(model, weights) -> Optional[DraftWeights]: + """The MTP layer's ``DraftWeights``: checkpoint layer ``num_hidden_layers``.""" + config = model.config + prefix = f"model.layers.{config.num_hidden_layers}" + eh_proj_key = f"{prefix}.eh_proj.weight" + if eh_proj_key not in weights: + return None + mapping = model.model_config.mapping + tp_size, tp_rank = mapping.tp_size, mapping.tp_rank + hidden = config.hidden_size + inter = config.moe_intermediate_size // tp_size + device = model.lm_head.weight.device + eh_proj = weights[eh_proj_key][:].to(device) + w3_w1 = torch.empty( + (config.n_routed_experts, 2 * inter, hidden), dtype=torch.bfloat16, device=device + ) + w2 = torch.empty((config.n_routed_experts, hidden, inter), dtype=torch.bfloat16, device=device) + + def shard(key: str, mode: TensorParallelMode) -> torch.Tensor: + tensor = load_weight_shard(weights[key], tp_size, tp_rank, mode, device=device) + assert tensor.dtype == torch.bfloat16, f"{key}: the MTP layer's experts are not BF16" + return tensor + + for expert in range(config.n_routed_experts): + name = f"{prefix}.mlp.experts.{expert}" + # The built-in MoE loader's slices: up_proj and gate_proj by rows, down_proj by columns. + w3_w1[expert, :inter].copy_(shard(f"{name}.up_proj.weight", TensorParallelMode.COLUMN)) + w3_w1[expert, inter:].copy_(shard(f"{name}.gate_proj.weight", TensorParallelMode.COLUMN)) + w2[expert].copy_(shard(f"{name}.down_proj.weight", TensorParallelMode.ROW)) + return DraftWeights(eh_proj, w3_w1, w2) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py new file mode 100644 index 000000000000..bec0b9684538 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Where a GlmMoeDsaForCausalLM config lands. Read this file and you know. + +``route`` is one forward-reading decision: the facts that fix which code runs, +compared against the deployments the targets below were written for. When none +fits it raises ``NoTarget`` stating those facts. In ``auto`` the engine then +uses the built-in implementation (``modeling_deepseekv3``); in ``require`` it +raises that error. +""" + +from __future__ import annotations + +from typing import Any, Optional, Tuple + +from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig + +from ..._router_index import NoTarget + +# (layers, hidden, routed experts, q_lora_rank). Shape-sniffing also matches +# fine-tunes of the same shape; the accuracy gate pins identity. +_SHAPE = (78, 6144, 256, 2048) + + +def route( + pretrained_config: Any, + mapping: Any, + compute_capability: Tuple[int, int], + spec_config: Optional[Any], +) -> type: + """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits. + + Each target is imported where it is returned, so declining costs no import + of modeling code. + """ + shape = ( + pretrained_config.num_hidden_layers, + pretrained_config.hidden_size, + pretrained_config.n_routed_experts, + pretrained_config.q_lora_rank, + ) + # tp16: attention heads and the expert intermediate dimension sharded over + # 16 ranks, no attention DP, no expert parallelism. A weight-layout + # question, which is why it selects a target rather than a runtime branch. + tp16 = ( + mapping.world_size == 16 + and mapping.tp_size == 16 + and mapping.moe_ep_size == 1 + and mapping.moe_tp_size == 16 + and not mapping.enable_attention_dp + and mapping.pp_size == 1 + and mapping.cp_size == 1 + ) + # The target drives the checkpoint's one MTP module through the one-model + # MTP-Eagle loop as a linear chain, at any draft length, or runs without + # speculative decoding. + mtp_eagle_or_none = spec_config is None or ( + isinstance(spec_config, MTPDecodingConfig) + and spec_config.spec_dec_mode.is_mtp_eagle_one_model() + and not getattr(spec_config, "use_dynamic_tree", False) + and getattr(spec_config, "eagle_choices", None) is None + and spec_config.max_total_draft_tokens == spec_config.max_draft_len + ) + # GLM-5.2-NVFP4 on 16 Rubin GPUs (sm107), tp16, MTP-Eagle or no speculation. + if compute_capability == (10, 7) and shape == _SHAPE and tp16 and mtp_eagle_or_none: + # One line on purpose: the claims suite reads the routed targets off + # `from ..modeling import ` lines. + # fmt: off + # isort: off + from .glm_5_2_nvfp4__sm_107__tp16.modeling import ModelingV2Glm52Nvfp4Sm107Tp16 + # isort: on + # fmt: on + return ModelingV2Glm52Nvfp4Sm107Tp16 + raise NoTarget( + f"compute capability {compute_capability}, config shape {shape}, " + f"world_size={mapping.world_size} tp={mapping.tp_size} moe_ep={mapping.moe_ep_size} " + f"moe_tp={mapping.moe_tp_size} attention_dp={mapping.enable_attention_dp} " + f"pp={mapping.pp_size} cp={mapping.cp_size}, speculative decoding " + f"{spec_config.spec_dec_mode.name if spec_config is not None else 'off'}; " + "the one glm_moe_dsa target is GLM-5.2-NVFP4 on sm107, tp16 (16 ranks, " + "tensor-parallel attention and experts, no attention DP), with one-model " + "MTP-Eagle as a linear chain or no speculative decoding" + ) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py index abe993e31438..bfd46a7f78e7 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py @@ -1248,7 +1248,9 @@ def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata): ) @staticmethod - def prepare(metadata: DSAtrtllmAttentionMetadata) -> None: + def prepare( + metadata: DSAtrtllmAttentionMetadata, *, defer_decode_outputs: bool = False + ) -> None: """ Prepare indexer for the forward pass. This should be called during metadata.prepare() stage. @@ -1256,6 +1258,9 @@ def prepare(metadata: DSAtrtllmAttentionMetadata) -> None: - Computes slot_mapping for KV cache updates - Prepares schedule_metadata for fp8_paged_mqa_logits - Stores generation request IDs for decode phase + + ``defer_decode_outputs`` leaves the slot mappings and schedules of a + generation-only batch to ``on_update_kv_lens``. """ # Skip indexer preparation if the kv_cache_manager doesn't have index_head_dim. @@ -1275,9 +1280,11 @@ def prepare(metadata: DSAtrtllmAttentionMetadata) -> None: ) # Store compressed KV token count for context requests metadata.num_ctx_kv_tokens = indexer_params.new_kv_tokens[:num_contexts].sum().item() + defer_decode_outputs = defer_decode_outputs and num_contexts == 0 # Prepare for update_k_cache - Indexer.prepare_for_update_k_cache(metadata, indexer_params) + if not defer_decode_outputs: + Indexer.prepare_for_update_k_cache(metadata, indexer_params) # Prepare for prefill phase if there are context requests if num_contexts > 0: @@ -1305,7 +1312,7 @@ def prepare(metadata: DSAtrtllmAttentionMetadata) -> None: Indexer.prepare_for_chunked_prefill(metadata, indexer_params) # Prepare for decode phase if there are generation requests - if num_generations > 0: + if num_generations > 0 and not defer_decode_outputs: # Prepare schedule metadata for fp8_paged_mqa_logits # This is a preprocessing step that computes scheduling information for the kernel Indexer.prepare_scheduler_metadata(metadata) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py index c6af4efd7dba..736d74e417af 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py @@ -6,15 +6,18 @@ import functools import os +from contextlib import contextmanager from dataclasses import dataclass, field from typing import TYPE_CHECKING, ClassVar, FrozenSet, List, Optional, Tuple import torch +from torch.utils.weak import WeakIdKeyDictionary import tensorrt_llm import tensorrt_llm.bindings from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE +from tensorrt_llm._torch.host_staging import copy_if_changed from tensorrt_llm._torch.utils import maybe_compile from tensorrt_llm._utils import get_sm_version, prefer_pinned from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata @@ -279,6 +282,11 @@ def __init__(self, *args, **kwargs): self.num_ctx_mla_kv_tokens = 0 self.nvfp4_mla_context_fp8_scratch = None self._radix_rows_per_sequence = 1 + # Set only inside defer_cuda_graph_decode_outputs(). + self._defer_decode_outputs = False + # Per INDEX_KEY host table, an object replaced whenever the table's + # contents change; device uploads of the table carry it as a token. + self._index_key_host_versions = WeakIdKeyDictionary() super().__init__(*args, **kwargs) sparse_metadata_params = self.sparse_metadata_params if not isinstance(sparse_metadata_params, DSAMetadataParams): @@ -357,10 +365,46 @@ def __post_init__(self): self.create_buffers_for_indexer(capture_graph=capture_graph) self._create_nvfp4_mla_generation_buffers(capture_graph=capture_graph) + @contextmanager + def defer_cuda_graph_decode_outputs(self): + """Engine-only scope for a selected decode graph's preparation. + + The engine must run ``on_update_kv_lens`` before any model consumer, + including warmup and capture. That call rebuilds the generation + indptrs, ``req_idx_per_token``, the indexer slot mappings, the DeepGEMM + schedules and the dense TopK indices from the KV lengths the forward + uses, so ``prepare`` leaves them to it for a generation-only batch. + Direct ``prepare`` callers retain fully initialized outputs. + """ + previous = getattr(self, "_defer_decode_outputs", False) + self._defer_decode_outputs = True + try: + yield + finally: + self._defer_decode_outputs = previous + + def _defers_decode_outputs(self) -> bool: + """Whether this ``prepare`` leaves the decode outputs to ``on_update_kv_lens``. + + Ragged verification, helix, a separate draft KV cache and non-native + DSA cache layouts keep the complete preparation. + """ + return ( + getattr(self, "_defer_decode_outputs", False) + and self.is_cuda_graph + and self.num_contexts == 0 + and self.num_generations > 0 + and not self.enable_ragged_verification + and not self.enable_helix + and self.draft_kv_cache_manager is None + and is_dsa_cache_manager(self.kv_cache_manager) + ) + def prepare(self): self._validate_ragged_verification() super().prepare() self._invalidate_pool_view_cache() + defer_decode_outputs = self._defers_decode_outputs() # Get kv lengths assert self.kv_cache_params.use_cache is True, "DSA requires use_cache to be True" @@ -382,13 +426,16 @@ def prepare(self): kv_lens = self._prepare_device_window_kv_lens(kv_lens) # For mla_rope_append_paged_kv_assign_q - self.prepare_for_mla_rope_append(cached_token_lens, kv_lens) + self.prepare_for_mla_rope_append( + cached_token_lens, kv_lens, generation_indptrs=not defer_decode_outputs + ) # Prepare to support skip indexer - self.prepare_for_skip_indexer(kv_lens) + self.prepare_for_skip_indexer(kv_lens, dense_topk_indices=not defer_decode_outputs) # For indices conversion - self.prepare_for_indices_conversion() + if not defer_decode_outputs: + self.prepare_for_indices_conversion() # For indexer k cache self.prepare_for_indexer_k_cache() @@ -397,7 +444,7 @@ def prepare(self): self.prepare_for_spec_decode(kv_lens) # Prepare metadata for indexer - Indexer.prepare(metadata=self) + Indexer.prepare(metadata=self, defer_decode_outputs=defer_decode_outputs) def _prepare_device_window_kv_lens(self, kv_lens: torch.Tensor) -> torch.Tensor: """Bound host-side KV extents until device-selected windows are installed.""" @@ -886,6 +933,10 @@ def create_buffers_for_mla_rope_append(self, capture_graph=False): device="cpu", pin_memory=prefer_pinned(), ) + # Entry 0 of an exclusive prefix sum is always 0, and the unfused + # rebuild in on_update_kv_lens() writes only the entries after it. + self.gen_cached_token_indptr[:1].zero_() + self.gen_kv_indptr[:1].zero_() def _create_nvfp4_mla_generation_buffers(self, capture_graph=False): """Allocate stable NVFP4 gather outputs for all MTP generation rows.""" @@ -1453,13 +1504,25 @@ def _update_indexer_k_cache_block_offsets(self) -> torch.Tensor: ) page_indices = pool_indices * cache_manager.indexer_k_cache_page_scale num_blocks = page_indices.shape[1] - self.host_indexer_k_cache_block_offsets[: self.num_seqs, :num_blocks].copy_(page_indices) - self.indexer_k_cache_block_offsets[: self.num_seqs].copy_( - self.host_indexer_k_cache_block_offsets[: self.num_seqs], non_blocking=True - ) - # Sanitize graph-padding entries that may be stale after cache + host_offsets = self.host_indexer_k_cache_block_offsets + host_rows = host_offsets[: self.num_seqs, :num_blocks] + version = self._index_key_host_versions.get(host_offsets) + if version is None or not torch.equal(host_rows, page_indices): + host_rows.copy_(page_indices) + version = self._index_key_host_versions[host_offsets] = object() + # Rows the device already took from this version of the host table + # are not uploaded again. + copy_if_changed( + self.indexer_k_cache_block_offsets[: self.num_seqs], + host_offsets[: self.num_seqs], + token=version, + ) + # The batch rows hold non-negative pool slots. Sanitize the + # graph-padding rows after them, which may be stale after cache # eviction or host-cache onboarding. - self.indexer_k_cache_block_offsets.clamp_(min=0) + padding_rows = self.indexer_k_cache_block_offsets[self.num_seqs :] + if padding_rows.numel() > 0: + padding_rows.clamp_(min=0) return pool_indices def set_skip_topk(self, skip: bool) -> None: @@ -2023,12 +2086,15 @@ def prepare_for_indexer_k_cache(self): host_block_table.masked_fill_( pad_cols.unsqueeze(0) >= num_blocks_per_seq[: self.num_seqs].unsqueeze(1), -1 ) - # Copy to GPU - self.block_table[: self.num_seqs, :max_blocks_used].copy_( - host_block_table, non_blocking=True - ) + # Copy to GPU unless the device rows already hold this table. + copy_if_changed(self.block_table[: self.num_seqs, :max_blocks_used], host_block_table) - def prepare_for_skip_indexer(self, kv_lens: torch.Tensor): + def prepare_for_skip_indexer(self, kv_lens: torch.Tensor, dense_topk_indices: bool = True): + """Decide which requests skip the indexer. + + ``dense_topk_indices=False`` leaves their dense TopK indices to + ``on_update_kv_lens``. + """ num_extra_kv_tokens = self.kv_cache_params.num_extra_kv_tokens if self.num_contexts > 0 and self.enable_indexer_skip: # Minus the number of extra KV tokens because when using one-model MTP, the @@ -2049,9 +2115,20 @@ def prepare_for_skip_indexer(self, kv_lens: torch.Tensor): ) else: self.skip_indexer_for_gen_reqs = False - self.prepare_dense_topk_indices(kv_lens) + if dense_topk_indices: + self.prepare_dense_topk_indices(kv_lens) - def prepare_for_mla_rope_append(self, cached_token_lens: torch.Tensor, kv_lens: torch.Tensor): + def prepare_for_mla_rope_append( + self, + cached_token_lens: torch.Tensor, + kv_lens: torch.Tensor, + generation_indptrs: bool = True, + ): + """Set the MLA RoPE-append lengths and indptrs. + + ``generation_indptrs=False`` leaves the generation indptrs to + ``on_update_kv_lens``. + """ if self.num_contexts > 0: self.num_ctx_cached_tokens = cached_token_lens[: self.num_contexts].sum().item() self.num_ctx_mla_kv_tokens = kv_lens[: self.num_contexts].sum().item() @@ -2085,6 +2162,8 @@ def prepare_for_mla_rope_append(self, cached_token_lens: torch.Tensor, kv_lens: if self.num_generations > 0: self.max_gen_seq_len = self.seq_lens[self.num_contexts : self.num_seqs].max().item() + if not generation_indptrs: + return # generation cached token indptr torch.cumsum( cached_token_lens[self.num_contexts : self.num_seqs], diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index d29def811bc5..1cbe4d393a8c 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -1077,20 +1077,30 @@ def prepare(self) -> None: # On the non-speculative path the host kv_lens snapshot bounds # every block-table access, so the staged/H2D width can be capped # at the batch's maximum instead of max_seq_len's worth of - # columns. Speculative decoding must stage the full width: - # draft/tree sub-steps and the overlap scheduler advance - # kv_lens_cuda on device past the host snapshot, and their - # kernels dereference block columns a host-derived cap would - # leave unstaged (uninitialized in this buffer). + # columns. Under speculative decoding the overlap scheduler and + # the draft sub-steps advance kv_lens_cuda on device past the host + # snapshot, and their kernels dereference block columns a + # host-derived cap would leave unstaged (uninitialized in this + # buffer). With a linear draft each of the two advances is less + # than one verification window of 1 + max_total_draft_tokens + # tokens, so the cap keeps two windows of room. Tree drafts and + # helix stage the full width. spec_active = (self.draft_kv_cache_manager is not None or self.is_spec_decoding_enabled or bool(self.kv_cache_params.num_extra_kv_tokens) or (self.runtime_features is not None and self.runtime_features.has_speculative_draft_tokens)) max_blocks = None - if not spec_active and self.kv_cache_manager.tokens_per_block: - max_blocks = ceil_div(max_kv_len, - self.kv_cache_manager.tokens_per_block) + tokens_per_block = self.kv_cache_manager.tokens_per_block + if tokens_per_block and not spec_active: + max_blocks = ceil_div(max_kv_len, tokens_per_block) + elif (tokens_per_block and self.max_total_draft_tokens is not None + and not self.is_spec_dec_tree + and not self.is_spec_dec_dynamic_tree + and not self.enable_helix): + max_blocks = ceil_div( + max_kv_len + 2 * (1 + self.max_total_draft_tokens), + tokens_per_block) self.kv_cache_manager.copy_batch_block_offsets( self.kv_cache_block_offsets, self.request_ids, diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index f500034c4df7..da37794ef6a7 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -349,6 +349,11 @@ def _(scores, scores_with_bias, n_group, topk_group, topk, dtype=scores_with_bias.dtype), scores.new_empty( shape, dtype=torch.int32) + @torch.library.register_fake("trtllm::glm_lowm_bf16_moe") + def _(x, router_logits, e_score_correction_bias, w3_w1, w2, top_k, + routed_scaling_factor): + return x.new_empty((x.shape[0], w2.shape[1])) + @torch.library.register_fake("trtllm::kimi_k3_noaux_tc_mxfp8_quant") def _(scores, bias, hidden_states, routed_scaling_factor): num_tokens = scores.shape[0] @@ -934,6 +939,35 @@ def _( logits.new_empty(logits.shape, dtype=torch.float32), ) + @torch.library.register_fake("trtllm::glm_verify_head_accept") + def _(target_tokens: torch.Tensor, draft_tokens: torch.Tensor, + num_contexts: int, accepted_tokens: Optional[torch.Tensor], + force_pool: Optional[torch.Tensor], + force_counter: Optional[torch.Tensor], force_mode: int, + force_base: int, force_max: int, force_frac: float, + counter_stride: int, slot_stride: int, pool_size: int): + return draft_tokens.new_empty((num_contexts + draft_tokens.shape[0], )) + + @torch.library.register_fake("trtllm::glm_draft_sample_glue") + def _(logits, exchange_buffer, exchange_flags, group_size, draft_tokens, + draft_step, embed_table, hidden, enorm_weight, hnorm_weight, eps, + col_lo, col_hi, glue_out, positions_in, positions_out, gather_ids, + seq_lens, fill_seq_lens, kv_lens, kv_update, rewind_accepted, + draft_len, block_offsets, dsa_layout, req_idx_per_token, + slot_mapping_fp8, slot_mapping_scale, gen_kv_indptr, + gen_cached_token_indptr, accepted_tokens, batch_indices, num_accepted, + next_new_tokens): + return None + + @torch.library.register_fake("trtllm::glm_draft_ehproj_table") + def _(logits, exchange_buffer, exchange_flags, group_size, draft_tokens, + draft_step, hidden_half, token_table, norm_weight, eps, normed_out, + residual_out, positions_in, positions_out, gather_ids, seq_lens, + fill_seq_lens, kv_lens, kv_update, rewind_accepted, draft_len, + block_offsets, dsa_layout, req_idx_per_token, slot_mapping_fp8, + slot_mapping_scale, gen_kv_indptr, gen_cached_token_indptr): + return None + @torch.library.register_fake("trtllm::fused_compute_probs_from_logits") def _( logits: torch.Tensor, @@ -985,6 +1019,39 @@ def _(pe: torch.Tensor, nope: torch.Tensor): scale = pe.new_empty((M, 1), dtype=torch.int32) return packed, scale + @torch.library.register_fake("trtllm::glm_fused_embedding") + def _(input: torch.Tensor, weight: torch.Tensor): + return weight.new_empty((*input.shape, weight.shape[1])) + + @torch.library.register_fake("trtllm::glm_dense_mlp_small_m") + def _(x, gate_up_weight, down_weight, gate_up_tiling=-1, down_tiling=-1): + return x.new_empty((x.shape[0], down_weight.shape[0])) + + @torch.library.register_fake("trtllm::glm_fused_a_gemm") + def _(hidden_states, weight, l2_prefetch): + return hidden_states.new_empty( + (hidden_states.shape[0], weight.shape[0])) + + @torch.library.register_fake("trtllm::glm_router_gemm") + def _(hidden_states, weight): + return hidden_states.new_empty( + (hidden_states.shape[0], weight.shape[0]), dtype=torch.float32) + + @torch.library.register_fake("trtllm::glm_mnnvl_hop") + def _(input, residual, gamma, epsilon, comm_buffer, buffer_flags, + shard_published): + return [torch.empty_like(residual), torch.empty_like(residual)] + + @torch.library.register_fake("trtllm::glm_mnnvl_moe_finalize_hop") + def _(gemm2_output, expert_weights, expanded_idx_to_permuted_idx, addend, + residual, gamma, epsilon, comm_buffer, buffer_flags): + return [torch.empty_like(residual), torch.empty_like(residual)] + + @torch.library.register_fake("trtllm::glm_mnnvl_allgather") + def _(input, comm_buffer, buffer_flags, group_size): + return input.new_empty((input.shape[0], input.shape[1] * group_size), + dtype=torch.float32) + @torch.library.register_fake("trtllm::fp8_quantize_1x128_packed_ue8m0") def _(input: torch.Tensor, use_r128c4_layout: bool = True) -> tuple[torch.Tensor, torch.Tensor]: diff --git a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py index 0063abfa4f1c..305b3e7912a0 100644 --- a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py @@ -461,7 +461,8 @@ def __init__(self, do_finalize: bool, act_type: int, tune_max_num_tokens: int = 8192, - use_dp: bool = False): + use_dp: bool = False, + prefer_gemm2_k128: bool = False): self.num_experts = num_experts self.top_k = top_k @@ -474,6 +475,7 @@ def __init__(self, self.routing_method_type = routing_method_type self.do_finalize = do_finalize self.act_type = act_type + self.prefer_gemm2_k128 = prefer_gemm2_k128 self.tuning_config = FP4BlockScaleMoERunner.get_tuning_config( self.num_experts // self.local_num_experts, @@ -483,8 +485,10 @@ def __init__(self, # The unique_id is used by the autotuner to get the cache key, so we hash on members # that influence tactic validity here. e.g. we are tuning FC1 and FC2 so the routing type does not matter def unique_id(self): - return (self.top_k, self.intermediate_size, self.local_num_experts, - self.act_type) + key = (self.top_k, self.intermediate_size, self.local_num_experts, + self.act_type) + # FC2 on gemm2K128 times the tile-8 tactics differently. + return key + ("gemm2_k128", ) if self.prefer_gemm2_k128 else key def get_runner(self): instance_key = (self.act_type, ) @@ -499,6 +503,10 @@ def forward( inputs: List[torch.Tensor], tactic: List[int] = [-1, -1], output: Optional[torch.Tensor] = None, + hidden_states_ready_event: Optional[int] = None, + router_hidden_states: Optional[torch.Tensor] = None, + router_weight: Optional[torch.Tensor] = None, + router_workspace: Optional[torch.Tensor] = None, ) -> torch.Tensor: assert isinstance(tactic, list) @@ -516,7 +524,9 @@ def forward( self.n_group, self.topk_group, self.intermediate_size, self.local_expert_offset, self.local_num_experts, self.routed_scaling_factor, self.routing_method_type, - self.do_finalize, tactic, args.topk_weights, args.topk_ids, output) + self.do_finalize, tactic, args.topk_weights, args.topk_ids, output, + hidden_states_ready_event, router_hidden_states, router_weight, + router_workspace, self.prefer_gemm2_k128) def get_valid_tactics(self, inputs: List[torch.Tensor], profile: OptimizationProfile, @@ -648,38 +658,61 @@ def get_tuning_config(cls, @torch.library.custom_op("trtllm::fp4_block_scale_moe_runner", mutates_args=()) -def fp4_block_scale_moe_runner(routing_logits: Optional[torch.Tensor], - routing_bias: Optional[torch.Tensor], - hidden_states: torch.Tensor, - hidden_states_scale: torch.Tensor, - gemm1_weights: torch.Tensor, - gemm1_weights_scale: torch.Tensor, - gemm1_bias: torch.Tensor, - gemm1_alpha: torch.Tensor, - gemm1_beta: torch.Tensor, - gemm1_clamp_limit: torch.Tensor, - gemm2_weights: torch.Tensor, - gemm2_weights_scale: torch.Tensor, - gemm2_bias: torch.Tensor, - output1_scale_scalar: torch.Tensor, - output1_scale_gate_scalar: torch.Tensor, - output2_scale_scalar: torch.Tensor, - num_experts: int, - top_k: int, - n_group: Optional[int], - topk_group: Optional[int], - intermediate_size: int, - local_expert_offset: int, - local_num_experts: int, - routed_scaling_factor: Optional[float], - routing_method_type: int, - do_finalize: bool, - act_type: int = ActType_TrtllmGen.SwiGlu.value, - topk_weights: Optional[torch.Tensor] = None, - topk_ids: Optional[torch.Tensor] = None, - output: Optional[torch.Tensor] = None, - tune_max_num_tokens: int = 8192, - use_dp: bool = False) -> List[torch.Tensor]: +def fp4_block_scale_moe_runner( + routing_logits: Optional[torch.Tensor], + routing_bias: Optional[torch.Tensor], + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm1_bias: torch.Tensor, + gemm1_alpha: torch.Tensor, + gemm1_beta: torch.Tensor, + gemm1_clamp_limit: torch.Tensor, + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + gemm2_bias: torch.Tensor, + output1_scale_scalar: torch.Tensor, + output1_scale_gate_scalar: torch.Tensor, + output2_scale_scalar: torch.Tensor, + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: Optional[float], + routing_method_type: int, + do_finalize: bool, + act_type: int = ActType_TrtllmGen.SwiGlu.value, + topk_weights: Optional[torch.Tensor] = None, + topk_ids: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, + tune_max_num_tokens: int = 8192, + use_dp: bool = False, + hidden_states_ready_event: Optional[int] = None, + router_hidden_states: Optional[torch.Tensor] = None, + router_weight: Optional[torch.Tensor] = None, + router_workspace: Optional[torch.Tensor] = None, + prefer_gemm2_k128: bool = False) -> List[torch.Tensor]: + """hidden_states_ready_event: raw cudaEvent_t (torch.cuda.Event.cuda_event) + recorded on another stream once hidden_states and hidden_states_scale are + written; the op waits on it right before FC1, after launching routing. + + router_hidden_states, router_weight, router_workspace: the router GEMV runs + in the routing kernel (DeepSeekV3 routing without groups, 256 experts, + top-8, hidden 6144, 1-8 tokens): the logits of + dsv3_router_gemm_op(router_hidden_states, router_weight.t()) are computed + there and routing_logits only gives their shape (it is not read or + written). router_workspace: int32 [FUSED_ROUTER_WORKSPACE_WORDS], see + make_fused_router_workspace; the kernel uses and restores it. Profiling + runs of the autotuner route from routing_logits as usual. + + prefer_gemm2_k128: FC2 runs on the persistent K = 128 kernel where it + covers the shape (NVFP4, tile N 8, intermediate size 128, no bias, enough + SMs), with the output rows of the Gemm2 cubin, bitwise; elsewhere the cubin + runs.""" tuner = AutoTuner.get() kernel_runner = FP4BlockScaleMoERunner( @@ -696,6 +729,7 @@ def fp4_block_scale_moe_runner(routing_logits: Optional[torch.Tensor], act_type, tune_max_num_tokens=tune_max_num_tokens, use_dp=use_dp, + prefer_gemm2_k128=prefer_gemm2_k128, ) # Prepare dummy topk tensors and hook for AutoTuner profiling @@ -755,7 +789,11 @@ def fp4_block_scale_moe_runner(routing_logits: Optional[torch.Tensor], result = kernel_runner( input_tensors, tactic=[-1, -1] if best_tactic == -1 else best_tactic, - output=output) + output=output, + hidden_states_ready_event=hidden_states_ready_event, + router_hidden_states=router_hidden_states, + router_weight=router_weight, + router_workspace=router_workspace) # When output is provided and do_finalize=True, the result is written in-place to output. # Return empty tensor to avoid aliasing constraint violation in PyTorch 2.9.1+ # (custom op output cannot be the same tensor as input). @@ -765,6 +803,24 @@ def fp4_block_scale_moe_runner(routing_logits: Optional[torch.Tensor], return result +# routingCustom::FusedRouterWorkspaceWords: two exchange buffers of +# [8 tokens][256 experts][4 slices] fp32, then the epoch word, padded to 16 B. +FUSED_ROUTER_WORKSPACE_WORDS = 2 * 8 * 256 * 4 + 4 + + +def make_fused_router_workspace(device: torch.device) -> torch.Tensor: + """Workspace of fp4_block_scale_moe_runner's fused router GEMV + routing + kernel (router_workspace): both exchange buffers idle (-0.0, bits + 0x80000000) and epoch 0. Between launches the buffer the next launch uses + is idle; nothing else may write the workspace. Launches that share a + workspace must run in stream order.""" + workspace = torch.zeros(FUSED_ROUTER_WORKSPACE_WORDS, + dtype=torch.int32, + device=device) + workspace[:FUSED_ROUTER_WORKSPACE_WORDS - 4].fill_(-2**31) + return workspace + + def fp4_block_scale_fake_output_without_finalize( hidden_states: Union[torch.Tensor, Fp4QuantizedTensor], num_experts: int, @@ -823,7 +879,12 @@ def _(routing_logits, topk_ids: Optional[torch.Tensor] = None, output: Optional[torch.Tensor] = None, tune_max_num_tokens: int = 8192, - use_dp: bool = False) -> List[torch.Tensor]: + use_dp: bool = False, + hidden_states_ready_event: Optional[int] = None, + router_hidden_states: Optional[torch.Tensor] = None, + router_weight: Optional[torch.Tensor] = None, + router_workspace: Optional[torch.Tensor] = None, + prefer_gemm2_k128: bool = False) -> List[torch.Tensor]: if do_finalize: num_tokens = hidden_states.shape[0] hidden_size = hidden_states.shape[1] * 2 diff --git a/tensorrt_llm/_torch/host_staging.py b/tensorrt_llm/_torch/host_staging.py index 38d25db83c7e..43e01119fb03 100644 --- a/tensorrt_llm/_torch/host_staging.py +++ b/tensorrt_llm/_torch/host_staging.py @@ -1,13 +1,75 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Bounded pinned staging for asynchronous host-to-device metadata uploads.""" +"""Asynchronous host-to-device metadata uploads. + +Bounded pinned staging, and skipping an upload whose payload the destination +already holds. +""" import torch +from torch.utils.weak import WeakIdKeyDictionary from tensorrt_llm._utils import maybe_pin_memory, prefer_pinned _STAGING_RING = 2 +# The last copy copy_if_changed() made into each storage, keyed by the +# storage's base tensor: views of one buffer share an entry, and the entry +# goes away with the buffer. +_last_copies = WeakIdKeyDictionary() + + +def _base(tensor: torch.Tensor) -> torch.Tensor: + return tensor if tensor._base is None else tensor._base + + +def _same_payload(recorded, current) -> bool: + if isinstance(recorded, torch.Tensor) or isinstance(current, torch.Tensor): + return ( + isinstance(recorded, torch.Tensor) + and isinstance(current, torch.Tensor) + and recorded.dtype == current.dtype + and recorded.shape == current.shape + and torch.equal(recorded, current) + ) + return recorded == current + + +def copy_if_changed(destination: torch.Tensor, values, *, token=None) -> bool: + """Copy ``values`` into ``destination`` unless it already holds them. + + ``values`` is a host tensor or a sequence of integers. ``destination`` + already holds them when the last copy made here into its storage wrote + the same elements with an equal payload. ``token``, when given, is + compared in place of the payload; it must change whenever the payload + does. A destination filtered here must not be written any other way + without a ``forget_copy`` call. Returns whether a copy was issued. + """ + if destination.numel() == 0: + return False + payload = values if isinstance(values, torch.Tensor) else tuple(values) + key = payload if token is None else token + region = (destination.data_ptr(), destination.dtype, destination.shape, destination.stride()) + base = _base(destination) + last = _last_copies.get(base) + if last is not None and last[0] == region and _same_payload(last[1], key): + return False + if not isinstance(payload, torch.Tensor): + payload = torch.tensor( + payload, dtype=destination.dtype, pin_memory=destination.is_cuda and prefer_pinned() + ) + destination.copy_(payload, non_blocking=destination.is_cuda) + _last_copies[base] = (region, key.clone() if isinstance(key, torch.Tensor) else key) + return True + + +def forget_copy(destination: torch.Tensor) -> None: + """Drop what ``copy_if_changed`` recorded for ``destination``'s storage. + + Call it after writing that storage any other way. + """ + _last_copies.pop(_base(destination), None) + def copy_host_to_device(store: dict, key, destination: torch.Tensor, value: torch.Tensor) -> None: """Asynchronously copy a host tensor into ``destination`` through pinned staging. diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index 6a54201f9065..b2bdc7cf78e7 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -1476,8 +1476,9 @@ def _get_decoder_layer_quant_config( def _compute_mlp_tp_size(self, intermediate_size: int, block_size: int) -> int: """ - For DeepSeek‑R1, MLP TP size is limited by intermediate_size // block_size - and must also be multiples of gpus_per_node to avoid expensive inter‑node allreduce. + For DeepSeek‑R1, MLP TP size is limited by intermediate_size // block_size. + Without the fused post-MLP all-reduce it is also capped at gpus_per_node to + avoid an expensive inter‑node allreduce. Args: intermediate_size (int): MLP intermediate size. @@ -1499,12 +1500,17 @@ def _compute_mlp_tp_size(self, intermediate_size: int, self.mapping.tp_size, ) - if tp > self.mapping.gpus_per_node: + if tp > self.mapping.gpus_per_node and not self.enable_fusion: + # Without the fused post-MLP all-reduce the MLP reduces within + # its own group, so a node-local MLP TP stays correct. mlp_tp_size = math.gcd( tp, self.mapping.gpus_per_node, ) # Avoid costly inter-node TP else: + # The fused post-MLP all-reduce sums over every rank of the TP + # group, so the MLP keeps the full TP: a smaller MLP TP would + # count each partial sum tp_size / mlp_tp_size times. mlp_tp_size = tp return mlp_tp_size diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index c74a0a8d9e7e..0357356d4d59 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -22,6 +22,7 @@ from tensorrt_llm._torch.compilation.piecewise_optimizer import PiecewiseRunner from tensorrt_llm._torch.compilation.utils import capture_piecewise_cuda_graph from tensorrt_llm._torch.distributed import Distributed +from tensorrt_llm._torch.host_staging import copy_if_changed from tensorrt_llm._torch.memory_buffer_utils import clear_memory_buffers, with_shared_pool from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.models.modeling_multimodal_mixin import _build_request_multimodal_input @@ -151,6 +152,28 @@ def _zero_fill_row_spans(buffer: torch.Tensor, spans: list[tuple[int, int]]) -> buffer[merged_begin:merged_end].fill_(0) +def _zero_rows_outside( + buffer: torch.Tensor, nonzero_rows: tuple[int, int], kept_rows: tuple[int, int] +) -> tuple[int, int]: + """Zero ``buffer`` everywhere except the half-open row span ``kept_rows``. + + ``buffer`` is zero outside the half-open span ``nonzero_rows``, so only + that span's rows outside ``kept_rows`` are cleared. Returns the span + outside which the buffer is zero once ``kept_rows`` is rewritten. + """ + begin, end = nonzero_rows + kept_begin, kept_end = kept_rows + if kept_begin >= kept_end: + if begin < end: + buffer[begin:end].zero_() + return (0, 0) + if begin < min(end, kept_begin): + buffer[begin : min(end, kept_begin)].zero_() + if max(begin, kept_end) < end: + buffer[max(begin, kept_end) : end].zero_() + return kept_rows + + def resolve_mamba_metadata_cls(model: torch.nn.Module) -> type[Mamba2Metadata]: """Resolve the model-specific Mamba metadata class with a default.""" return getattr(model, "mamba_metadata_cls", None) or Mamba2Metadata @@ -470,9 +493,6 @@ def _allocate_decoder_buffers(self) -> None: self.gather_ids_cuda = torch.empty( (self._config.max_num_tokens,), dtype=torch.int, device="cuda" ) - self.num_accepted_draft_tokens_cuda = torch.empty( - (self._config.max_batch_size,), dtype=torch.int, device="cuda" - ) self.previous_pos_indices_cuda = torch.empty( (self._config.max_num_tokens,), dtype=torch.int, device="cuda" ) @@ -482,6 +502,9 @@ def _allocate_decoder_buffers(self) -> None: self.previous_kv_lens_offsets_cuda = torch.zeros( (self._config.max_batch_size,), dtype=torch.int, device="cuda" ) + # Half-open row spans outside which the two offset buffers are zero. + self._previous_pos_id_offsets_rows = (0, 0) + self._previous_kv_lens_offsets_rows = (0, 0) self.previous_batch_indices_cuda = torch.empty( (self._config.max_num_tokens,), dtype=torch.int, device="cuda" ) @@ -2856,8 +2879,16 @@ def _preprocess_inputs( Make some changes to the device inputs and avoid blocking the async data transfer """ attn_meta = inputs.get("attn_metadata") + # The overlap correction below changes kv_lens_cuda and refreshes the metadata after it. + refreshed_after_correction = ( + enable_spec_decode + and not self._config.disable_overlap_scheduler + and attn_meta is not None + and attn_meta.kv_cache_manager is not None + and hasattr(attn_meta, "kv_lens_cuda") + ) # Invalidate per-forward-pass caches so they are recomputed (and captured) on every _forward_step. - if attn_meta is not None: + if attn_meta is not None and not refreshed_after_correction: attn_meta.on_update_kv_lens() if enable_spec_decode and not self._config.disable_overlap_scheduler: @@ -3319,6 +3350,22 @@ def _build_multimodal_data_for_llm( return request.py_multimodal_data return self._mm_item_scheduler.build_multimodal_data_for_llm(request, runtime) + def _zero_previous_offsets_outside( + self, pos_id_rows: tuple[int, int], kv_lens_rows: tuple[int, int] + ) -> None: + """Zero the overlap position-id and KV-length offsets outside the + half-open row spans the caller rewrites next. + + Only these writes change the two buffers, so rows outside the spans + rewritten last are already zero and are not cleared again. + """ + self._previous_pos_id_offsets_rows = _zero_rows_outside( + self.previous_pos_id_offsets_cuda, self._previous_pos_id_offsets_rows, pos_id_rows + ) + self._previous_kv_lens_offsets_rows = _zero_rows_outside( + self.previous_kv_lens_offsets_cuda, self._previous_kv_lens_offsets_rows, kv_lens_rows + ) + def _prepare_tp_inputs( self, scheduled_requests: ScheduledRequests, @@ -3421,7 +3468,6 @@ def _prepare_tp_inputs( mrope_dummy_seq_slot = get_mrope_dummy_seq_slot( self._config.max_num_tokens, self.mapping.pp_size ) - num_accepted_draft_tokens = [] # per request extra_inputs = self._new_extra_inputs() context_prompt_lookahead = None @@ -3458,7 +3504,6 @@ def _prepare_tp_inputs( gather_ids.append(len(input_ids) - 1) sequence_lengths.append(len(prompt_tokens)) - num_accepted_draft_tokens.append(len(prompt_tokens) - 1) prompt_lengths.append(len(prompt_tokens)) past_seen_token_num = begin_compute num_cached_tokens_per_seq.append(past_seen_token_num - request.py_num_compressed_tokens) @@ -3721,7 +3766,6 @@ def _helix_pack_extend(request, group: int) -> int: prompt_lengths.append(request.py_prompt_len) sequence_lengths.append(1 + num_draft_tokens) - num_accepted_draft_tokens.append(num_draft_tokens) gather_ids.extend( list(range(len(position_ids), len(position_ids) + 1 + num_draft_tokens)) ) @@ -3761,7 +3805,6 @@ def _helix_pack_extend(request, group: int) -> int: request.py_batch_idx = request.py_seq_slot sequence_lengths.append(runtime_tokens_per_gen_step) - num_accepted_draft_tokens.append(request.py_num_accepted_draft_tokens) past_seen_token_num = request.max_beam_num_tokens - 1 draft_lens.append(runtime_draft_token_buffer_width) @@ -3826,7 +3869,6 @@ def _helix_pack_extend(request, group: int) -> int: - 1 - (self._config.original_max_draft_len - request.py_num_accepted_draft_tokens) ) - num_accepted_draft_tokens.append(request.py_num_accepted_draft_tokens) sequence_lengths.append(1 + self._config.original_max_draft_len) prompt_lengths.append(request.py_prompt_len) @@ -3889,10 +3931,9 @@ def _helix_pack_extend(request, group: int) -> int: ) # Pre-extend constant-value lists to avoid per-request append - # overhead (saves ~3 append calls per request). + # overhead (saves ~2 append calls per request). draft_lens.extend([0] * (_n_gen * beam_width)) sequence_lengths.extend([1] * (_n_gen * beam_width)) - num_accepted_draft_tokens.extend([0] * (_n_gen * beam_width)) for request in generation_requests: request_ids.append(request.py_request_id) @@ -4058,11 +4099,8 @@ def _helix_pack_extend(request, group: int) -> int: previous_batch_len = len(previous_batch_indices) def previous_seq_slots_device(): - previous_batch_indices_host = torch.tensor( - previous_batch_indices, dtype=torch.int, pin_memory=prefer_pinned() - ) previous_slots = self.previous_batch_indices_cuda[:previous_batch_len] - previous_slots.copy_(previous_batch_indices_host, non_blocking=True) + copy_if_changed(previous_slots, previous_batch_indices) return previous_slots if scheduled_requests.is_attention_dp_phase_idle: @@ -4083,19 +4121,24 @@ def previous_seq_slots_device(): if num_draft_tokens > 0: draft_tokens = torch.tensor(draft_tokens, dtype=torch.int, pin_memory=prefer_pinned()) self.draft_tokens_cuda[: len(draft_tokens)].copy_(draft_tokens, non_blocking=True) - if self._config.is_spec_decode and len(num_accepted_draft_tokens) > 0: - num_accepted_draft_tokens = torch.tensor( - num_accepted_draft_tokens, dtype=torch.int, pin_memory=prefer_pinned() - ) - self.num_accepted_draft_tokens_cuda[: len(num_accepted_draft_tokens)].copy_( - num_accepted_draft_tokens, non_blocking=True - ) if next_draft_tokens_device is not None: - # Initialize these two values to zeros - self.previous_pos_id_offsets_cuda *= 0 - self.previous_kv_lens_offsets_cuda *= 0 runtime_tokens_per_gen_step = self.get_runtime_tokens_per_gen_step(runtime_draft_len) runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 + # Initialize these two values to zeros, except the rows of the + # previous-batch requests, which are written below. + num_extend_reqeust_wo_dummy = len(extend_requests) - len(extend_dummy_requests) + previous_request_rows = ( + (num_extend_reqeust_wo_dummy - previous_batch_len, num_extend_reqeust_wo_dummy) + if previous_batch_len > 0 + else (0, 0) + ) + self._zero_previous_offsets_outside( + ( + previous_request_rows[0] * runtime_tokens_per_gen_step, + previous_request_rows[1] * runtime_tokens_per_gen_step, + ), + previous_request_rows, + ) if previous_batch_len > 0: previous_slots = previous_seq_slots_device() @@ -4121,11 +4164,8 @@ def previous_seq_slots_device(): ) # prepare data for the preprocess inputs kv_len_offsets_device = new_tokens_lens_device - runtime_tokens_per_gen_step - previous_pos_indices_host = torch.tensor( - previous_pos_indices, dtype=torch.int, pin_memory=prefer_pinned() - ) - self.previous_pos_indices_cuda[0:previous_batch_tokens].copy_( - previous_pos_indices_host, non_blocking=True + copy_if_changed( + self.previous_pos_indices_cuda[0:previous_batch_tokens], previous_pos_indices ) # The order of requests in a batch: [context requests, generation requests] @@ -4148,7 +4188,6 @@ def previous_seq_slots_device(): # For 3) 'dummy requests': pad dummy requests for CUDA graph or attention dp. # Already set to '0' during initialization. - num_extend_reqeust_wo_dummy = len(extend_requests) - len(extend_dummy_requests) self.previous_pos_id_offsets_cuda[ (num_extend_reqeust_wo_dummy - previous_batch_len) * runtime_tokens_per_gen_step : num_extend_reqeust_wo_dummy @@ -4181,8 +4220,7 @@ def previous_seq_slots_device(): # so we need to set the previous_pos_id_offsets and previous_kv_lens_offsets to zeros # to skip the value changes in _preprocess_inputs. Otherwise, there will be illegal memory access # when writing key/values to the KV cache. - self.previous_pos_id_offsets_cuda *= 0 - self.previous_kv_lens_offsets_cuda *= 0 + self._zero_previous_offsets_outside((0, 0), (0, 0)) # The host copy covers input_ids_cuda[:num_tokens] and the device copies # cover the previous-batch rows; padding-request rows are covered by @@ -4234,10 +4272,7 @@ def previous_seq_slots_device(): final_position_ids = self.position_ids_cuda[:total_num_tokens].unsqueeze(0) if enable_spec_decode: - self.gather_ids_cuda[: len(gather_ids)].copy_( - torch.tensor(gather_ids, dtype=torch.int, pin_memory=prefer_pinned()), - non_blocking=True, - ) + copy_if_changed(self.gather_ids_cuda[: len(gather_ids)], gather_ids) if self.mapping.has_cp_helix(): # A non-None owned-count list is what arms @@ -4479,9 +4514,6 @@ def previous_seq_slots_device(): # num_generations / num_tokens / seq_lens are set above, before the # attention-DP allgather that must agree with prepare(). spec_metadata.host_position_ids = host_position_ids - spec_metadata.num_accepted_draft_tokens = self.num_accepted_draft_tokens_cuda[ - : len(num_accepted_draft_tokens) - ] if context_prompt_lookahead is not None: spec_metadata.populate_context_prompt_lookahead(context_prompt_lookahead) # No-op for non 1-model @@ -5246,8 +5278,13 @@ def _cuda_graph_metadata_prepare_scope( """Allow derived outputs only when this forward owns a decode graph. Both capture/warmup and replay execute _preprocess_inputs before model - consumers. Non-speculative overlap has the same ordering; speculative - acceptance corrections and draft-manager swaps keep full preparation. + consumers, and it runs on_update_kv_lens, which rebuilds the derived + outputs. Non-speculative overlap has the same ordering. When a linear + one-model speculative step verifies drafts, that call follows the + overlap acceptance correction of kv_lens_cuda, so the outputs come + from the corrected lengths. Draft-manager swaps and ragged + verification rebuild metadata outside the graph and keep full + preparation, as does a speculative step without drafts. """ defer = getattr(attn_metadata, "defer_cuda_graph_decode_outputs", None) if ( @@ -5255,11 +5292,14 @@ def _cuda_graph_metadata_prepare_scope( and isinstance(key, KeyType) and attn_metadata.is_cuda_graph is True and key.num_contexts == 0 - and key.draft_len == 0 + and ( + key.draft_len > 0 and self._linear_one_model_spec_decode(attn_metadata) + if enable_spec_decode + else key.draft_len == 0 + ) and scheduled_requests.num_context_requests == 0 and scheduled_requests.num_generation_requests > 0 and not promoted_context_request_ids - and not enable_spec_decode and not self.use_beam_search and self._context_graph_promotion_supported and self.mapping.cp_size == 1 @@ -5267,6 +5307,18 @@ def _cuda_graph_metadata_prepare_scope( return defer() return contextlib.nullcontext() + def _linear_one_model_spec_decode(self, attn_metadata: AttentionMetadata) -> bool: + """Whether speculation verifies a linear draft in the target model's + own forward, with no separate draft KV cache and no ragged windows.""" + spec_config = self._config.spec_config + return ( + spec_config is not None + and spec_config.spec_dec_mode.use_one_engine() + and spec_config.is_linear_tree + and getattr(attn_metadata, "draft_kv_cache_manager", None) is None + and not getattr(attn_metadata, "enable_ragged_verification", False) + ) + def _get_remote_tail_decode_rank_tokens( self, attn_metadata: AttentionMetadata, can_run_graph: bool ) -> list[int]: diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder_decoder.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder_decoder.py index 8178deb67fea..60e3177295fb 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder_decoder.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder_decoder.py @@ -17,6 +17,7 @@ from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.attention.backends.vanilla import VanillaAttentionMetadata from tensorrt_llm._torch.distributed import Distributed +from tensorrt_llm._torch.host_staging import forget_copy from tensorrt_llm._torch.memory_buffer_utils import with_shared_pool from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import MoeLoadBalancer @@ -1350,6 +1351,8 @@ def _prepare_encoder_decoder_inputs_fast( buffers["previous_batch_indices"][:num_previous_batch_requests], non_blocking=True, ) + # The decoder staging writes this buffer through copy_if_changed(). + forget_copy(previous_slots) self._encoder_decoder_staged_request_ids = staged_request_ids generation_begin = num_context_tokens generation_end = generation_begin + num_previous_batch_requests diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 9aa7d619a3e8..f1aad6acd9da 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -6536,8 +6536,6 @@ def copy_batch_block_offsets( "Offloaded sparse history requires Batch metadata and sparse fetch; " "dense attention offsets cannot address host slots" ) - # max_blocks is accepted for signature parity with KVCacheManager; the - # device-side copy op here already scales with allocated blocks only. assert beam_width <= self.max_copy_beam_width # Cross KV is request-scoped: generation beams share the encoder K/V @@ -6556,6 +6554,14 @@ def copy_batch_block_offsets( ) return + # As in KVCacheManager, max_blocks bounds the copied block columns of + # each row; the columns past it keep their previous contents. Helix + # rows are copied whole. + bound = ( + {"max_num_blocks": max_blocks} + if max_blocks is not None and not self._has_cp_helix + else {} + ) copy_batch_block_offsets_to_device( self.host_kv_cache_block_offsets, dst_tensor, @@ -6563,6 +6569,7 @@ def copy_batch_block_offsets( self.index_scales, self.kv_offset, self._stream.cuda_stream, + **bound, ) @staticmethod diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index fc29a0633f41..bfc7ccba0a7f 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -16,6 +16,7 @@ from ..attention.backends import AttentionMetadata from ..attention.backends.flashinfer import FlashInferAttentionMetadata from ..attention.backends.sparse.params import MTPIndexShareMetadata +from ..host_staging import copy_if_changed from ..model_config import ModelConfig from ..pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.llm_request import LlmRequest @@ -348,12 +349,7 @@ def prepare(self): assert self.request_ids is not None # update batch indices num_seqs = len(self.request_ids) - batch_indices = torch.arange(num_seqs, - dtype=torch.int, - device='cpu', - pin_memory=prefer_pinned()) - self.batch_indices_cuda[:num_seqs].copy_(batch_indices, - non_blocking=True) + copy_if_changed(self.batch_indices_cuda[:num_seqs], range(num_seqs)) # Snapshot the real (unpadded) scheduled token count before the decrement # below. maybe_capture_hidden_states runs during the target forward and # must bound the capture by this count; self.num_tokens is about to be diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 8237c09ee8b7..ba1d1b8b8ad8 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -836,7 +836,8 @@ def _populate_request_rng_state( Both layouts are produced: ``request_*`` with one entry per request, and ``seeds`` / ``offsets`` expanded to one entry per logits row (the temperatures / top_ks / top_ps / min_ps layout), because the sampling calls - take one or the other. + take one or the other. Only rejection sampling and its draft sampling + read ``request_*``, so without it those buffers are not uploaded. A request that specified no seed gets ``DEFAULT_SAMPLING_SEED`` and an offset window from the counter shared by all unseeded requests. Each of @@ -899,8 +900,9 @@ def _upload(dst: torch.Tensor, values: list[int]) -> None: pin_memory=prefer_pinned()), non_blocking=True) - _upload(self.request_seeds, request_seeds) - _upload(self.request_offsets, request_offsets) + if self.use_rejection_sampling: + _upload(self.request_seeds, request_seeds) + _upload(self.request_offsets, request_offsets) _upload(self.seeds, flat_seeds) _upload(self.offsets, flat_offsets) diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index 5fb5b185dc0b..02688203d157 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -28,6 +28,7 @@ import torch from ...sampling_params import SamplingParams +from ..host_staging import copy_if_changed from ..pyexecutor.llm_request import LlmRequest, LlmRequestState, get_draft_token_length from ..pyexecutor.resource_manager import BaseResourceManager from ..pyexecutor.sampler import ( @@ -166,6 +167,8 @@ class SpecSampler(Sampler[SampleStateSpec], AsyncWorkerMixin): """ SampleState = SampleStateSpec + # Store rows of the sampled requests in batch order; allocated on first use. + _sampling_slots: Optional[torch.Tensor] = None def is_generation_model(self) -> bool: return True @@ -442,8 +445,10 @@ def sample_async( for r in sampling_requests ] - slots = torch.as_tensor([r.py_seq_slot for r in sampling_requests], dtype=torch.long) - slots = slots.to(device="cuda", non_blocking=True) + if self._sampling_slots is None: + self._sampling_slots = torch.empty_like(self.store.new_tokens_lens, dtype=torch.long) + slots = self._sampling_slots[:num_sampling_requests] + copy_if_changed(slots, [r.py_seq_slot for r in sampling_requests]) o_new_tokens = outputs["new_tokens"][num_skip : num_skip + num_sampling_requests] o_new_tokens_lens = outputs["new_tokens_lens"][num_skip : num_skip + num_sampling_requests] diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_mla_kv_half_append.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_mla_kv_half_append.py new file mode 100644 index 000000000000..c3ad37ea9d02 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_mla_kv_half_append.py @@ -0,0 +1,288 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM-5.2 tp16 mla_kv_half_append entry. + +The paged state is real: a DSACacheManagerV2 over GLM-5.2's latent rows (FP8, 576 per token, +64-token pages) with several layers, the attention metadata of one generation step after a +context, and the MLA rope table at GLM-5.2's rope_theta. The rows the op appends are read back +where main's `convert_req_index_to_global` places their positions -- the rows main's built-in +sparse MLA decode reads -- and gated against the entry's reference; no other byte of the pool +may change. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.attention.mla_kv_half_append import ( # noqa: E501 + KV_LORA_RANK, + LATENT_OFFSET, + ROPE_DIM, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.attention.mla_kv_half_append import ( # noqa: E501 + mla_kv_half_append as op, +) +from tensorrt_llm._torch.attention.backends.interface import RopeParams +from tensorrt_llm._torch.attention.backends.sparse.dsa.cache_manager import DSACacheManagerV2 +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig, KvCacheConfig +from tensorrt_llm.mapping import Mapping + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or get_sm_version() != 107, + reason="the GLM-5.2 tp16 target's ops are certified on sm_107", +) + +HEAD_DIM = KV_LORA_RANK + ROPE_DIM +MAX_SEQ_LEN = 8192 +#: One E4M3 step: 2**-3 relative at the bottom of a binade, 2**-9 among the subnormals. +_E4M3_STEP_RTOL = 2**-3 +_E4M3_SUBNORMAL_STEP = 2**-9 +#: How much of an appended row may miss bit-exactness. See `_compare`. +_MAX_INEXACT_FRACTION = 1e-3 + + +def _compare(got: torch.Tensor, want: torch.Tensor) -> None: + """One E4M3 step, and at most 0.1 % of the bytes off bit-exactness. + + Kernel and reference evaluate the same fp32 expressions and round once; they differ in the + order of the 512-term sum of squares, in rsqrtf against torch.rsqrt (2 ulp), and in FMA + contraction of the rotation -- about 1e-7 relative, which crosses an E4M3 rounding boundary + (3 mantissa bits) for about one value in a million, and then by one step. A reference that + rounds the normalized latent and the rotated K to BF16 first (mla_rope_generation's KV + half) misses bit-exactness on about 3 % of the bytes, so the floor tells the two apart; a + wrong weight, epsilon, pair layout or position moves most bytes by several steps. + """ + assert got.shape == want.shape, (got.shape, want.shape) + torch.testing.assert_close( + got.float(), want.float(), rtol=_E4M3_STEP_RTOL, atol=_E4M3_SUBNORMAL_STEP + ) + inexact = int((got.view(torch.uint8) != want.view(torch.uint8)).sum()) + allowed = max(1, int(_MAX_INEXACT_FRACTION * got.numel())) + assert inexact <= allowed, f"{inexact}/{got.numel()} bytes not bit-exact (allowed {allowed})" + + +def _fill_e4m3_(dst: torch.Tensor, gen: torch.Generator) -> None: + """Fill a contiguous uint8 view with E4M3 values of N(0, 1).""" + flat = dst.reshape(-1) + chunk = 1 << 26 + for start in range(0, flat.numel(), chunk): + stop = min(start + chunk, flat.numel()) + values = torch.randn(stop - start, generator=gen, device=flat.device) + flat[start:stop].copy_(values.to(torch.float8_e4m3fn).view(torch.uint8)) + + +class _DsaEnv: + """A DSACacheManagerV2 over GLM-5.2's latent rows and one generation step's metadata.""" + + def __init__(self, num_layers: int, tokens_per_block: int, rope_theta: float) -> None: + self.tokens_per_block = tokens_per_block + self.max_batch_size = 2 + self.manager = DSACacheManagerV2( + kv_cache_config=KvCacheConfig( + max_tokens=self.max_batch_size * (MAX_SEQ_LEN + tokens_per_block), + enable_block_reuse=False, + ), + kv_cache_type=CacheType.SELFKONLY, + num_layers=num_layers, + num_kv_heads=1, + head_dim=HEAD_DIM, + tokens_per_block=tokens_per_block, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=self.max_batch_size, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=DataType.FP8, + sparse_attention_config=DeepSeekSparseAttentionConfig( + index_head_dim=128, index_n_heads=32, index_topk=2048 + ), + ) + rope = RopeParams( + dim=ROPE_DIM, theta=rope_theta, max_positions=MAX_SEQ_LEN, duplicate_data=True + ) + _, self.rotary_cos_sin = rope.create_rope_const_params() + + def generation_step(self, cached: list, seq_len_q: int) -> TrtllmAttentionMetadata: + """Requests 0 .. n-1 with `cached` tokens each, grown by one step of `seq_len_q`.""" + request_ids = list(range(len(cached))) + self.manager.add_dummy_requests(request_ids, token_nums=cached) + for rid in request_ids: + cache = self.manager.kv_cache_map[rid] + assert cache.resize(cache.capacity + seq_len_q) + metadata = TrtllmAttentionMetadata( + max_num_requests=self.max_batch_size, + max_num_tokens=8192, + kv_cache_manager=self.manager, + ) + metadata.seq_lens = torch.tensor([seq_len_q] * len(cached), dtype=torch.int) + metadata.num_contexts = 0 + metadata.request_ids = request_ids + metadata.prompt_lens = cached + metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=cached) + metadata.prepare() + return metadata + + def pool_rows(self) -> torch.Tensor: + """The pool as [rows, 576] bytes, the rows convert_req_index_to_global counts.""" + pool = self.manager.get_unique_primary_pool() + return pool.view(torch.uint8).reshape(-1, HEAD_DIM) + + def global_rows(self, kv_lens: list, positions: torch.Tensor, seq_len_q: int, layer: int): + """Main's pool rows of each generation token's position in `layer`.""" + tpb = self.tokens_per_block + slots = self.manager.get_pool_block_indices( + len(kv_lens), request_ids=list(range(len(kv_lens))), num_contexts=0 + ) + blocks = torch.tensor([-(-n // tpb) for n in kv_lens]) + width = int(blocks.max()) + table = slots[:, :width].to(torch.int32).clone() + table[torch.arange(width)[None, :] >= blocks[:, None]] = -1 + page_index_scale, layer_offset = self.manager.get_primary_pool_page_index_params(layer) + req_idx = (torch.arange(positions.shape[0], device="cuda") // seq_len_q).to(torch.int32) + rows = torch.ops.trtllm.convert_req_index_to_global( + req_idx, + table.cuda(), + positions.to(torch.int32)[:, None].contiguous(), + tpb, + 1, + page_index_scale * tpb, + layer_offset, + ) + return rows[:, 0].long() + + def shutdown(self) -> None: + self.manager.shutdown() + + +def _step_args(env: _DsaEnv, spec: dict, gen: torch.Generator) -> tuple: + """One generation step's call: the metadata, a fused-A output and the scale buffers.""" + cached = list(spec["cached"]) + seq_len_q = spec["seq_len_q"] + metadata = env.generation_step(cached, seq_len_q) + num_tokens = len(cached) * seq_len_q + fused_a = torch.randn(num_tokens, spec["latent_width"], generator=gen, device="cuda").to( + torch.bfloat16 + ) + kv_norm_weight = (1.0 + 0.1 * torch.randn(KV_LORA_RANK, generator=gen, device="cuda")).to( + torch.bfloat16 + ) + return ( + fused_a[:, LATENT_OFFSET : LATENT_OFFSET + HEAD_DIM], + kv_norm_weight, + spec["kv_norm_eps"], + env.rotary_cos_sin, + metadata.kv_lens_cuda_runtime, + 0, + seq_len_q, + metadata.kv_cache_block_offsets, + env.manager.kv_cache_pool_pointers, + env.manager.kv_cache_pool_mapping, + spec["layer_idx"], + spec["tokens_per_block"], + env.manager.max_seq_len, + torch.full((2,), float("nan"), dtype=torch.float32, device="cuda"), + torch.full((1,), float("nan"), dtype=torch.float32, device="cuda"), + spec["softmax_scale"], + spec["trigger_at_entry"], + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by name + certifying(op) + spec = cell.spec + seed = zlib.crc32(cell.why.encode()) % 2**31 + gen = torch.Generator(device="cuda").manual_seed(seed) + env = _DsaEnv(spec["num_layers"], spec["tokens_per_block"], spec["rope_theta"]) + try: + args = _step_args(env, spec, gen) + latent, bmm1, bmm2 = args[0], args[13], args[14] + pool_rows = env.pool_rows() + _fill_e4m3_(pool_rows, gen) + pool_before = pool_rows.clone() + latent_before = latent.clone() + + op(*args) + torch.cuda.synchronize() + want_rows, positions, want_bmm1, want_bmm2 = op.reference(*args) + + # The position of each row: this step's tokens are the last seq_len_q of each sequence. + seq_len_q = spec["seq_len_q"] + kv_lens = [c + seq_len_q for c in spec["cached"]] + assert positions.tolist() == [n - seq_len_q + t for n in kv_lens for t in range(seq_len_q)] + + # Each row lands where main's layout puts its position in this layer. + rows = env.global_rows(kv_lens, positions, seq_len_q, spec["layer_idx"]) + assert rows.unique().numel() == rows.numel() and bool((rows >= 0).all()) + _compare(pool_rows[rows].view(torch.float8_e4m3fn), want_rows) + + # And nothing else moves: not another layer, slot or position, not the latent. + changed = (pool_rows != pool_before).any(dim=-1).nonzero().flatten() + assert set(changed.tolist()) <= set(rows.tolist()), "the op wrote outside its rows" + assert torch.equal(latent, latent_before) + + # The attention scales, bit for bit: one fp32 product. + assert torch.equal(bmm1, want_bmm1), (bmm1, want_bmm1) + assert torch.equal(bmm2, want_bmm2), (bmm2, want_bmm2) + + # Rerunning the prepared step rewrites the same bytes. + appended = pool_rows[rows].clone() + op(*args) + torch.cuda.synchronize() + assert torch.equal(pool_rows[rows], appended) + finally: + env.shutdown() + + +def _guard_args(**overrides) -> tuple: + """The first cell's call on a fresh environment, with arguments replaced by position.""" + spec = op.CELLS[0].spec + gen = torch.Generator(device="cuda").manual_seed(0) + env = _DsaEnv(spec["num_layers"], spec["tokens_per_block"], spec["rope_theta"]) + args = list(_step_args(env, spec, gen)) + names = ("sequence_length", "rotary_cos_sin", "kv_cache_block_offsets") + positions = {"sequence_length": 4, "rotary_cos_sin": 3, "kv_cache_block_offsets": 7} + for name, value in overrides.items(): + assert name in names + args[positions[name]] = value(args[positions[name]]) + return env, tuple(args) + + +def test_guard_refuses_a_sequence_shorter_than_its_step() -> None: + """Rows at negative positions are skipped: those tokens never reach the cache.""" + env, args = _guard_args(sequence_length=lambda t: torch.full_like(t, 5)) + try: + with pytest.raises(AssertionError, match="unwritten"): + with validating(op): + op(*args) + finally: + env.shutdown() + + +def test_guard_refuses_a_rope_table_shorter_than_the_positions() -> None: + """The table is read at the position without a bound.""" + env, args = _guard_args(rotary_cos_sin=lambda t: t.reshape(-1)[: 100 * 2 * ROPE_DIM]) + try: + with pytest.raises(AssertionError, match="reads past it"): + with validating(op): + op(*args) + finally: + env.shutdown() + + +def test_guard_refuses_block_offsets_narrower_than_the_positions() -> None: + """A block past the offsets' width is read from the next sequence's offsets.""" + env, args = _guard_args(kv_cache_block_offsets=lambda t: t[..., :8].contiguous()) + try: + with pytest.raises(AssertionError, match="needs more"): + with validating(op): + op(*args) + finally: + env.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_sparse_mla_flash_decode.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_sparse_mla_flash_decode.py new file mode 100644 index 000000000000..a255b9e9819b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_sparse_mla_flash_decode.py @@ -0,0 +1,559 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM-5.2 tp16 sparse_mla_flash_decode entry. + +The cells build GLM-5.2's TP16 pool geometry directly: FP8 latent rows in 79 layers' 64-token +pages per pool slot, a block table that maps each sequence's blocks to shuffled slots (-1 past +the sequence), and top-k rows laid out as the DSA indexer writes them -- the valid entries +packed in front, this step's visible tokens among them, -1 behind. Every cell is driven through +the entry's guard and gated against the entry's reference. + +Two further tests: a captured CUDA graph replays the shared-layer cell bit for bit, and on a +real DSACacheManagerV2 the KV append writes this step's rows where main's +`convert_req_index_to_global` says they live and the flash decode's in-kernel remap reads +exactly those rows. +""" + +import math +import zlib +from typing import NamedTuple + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.attention.mla_kv_half_append import ( # noqa: E501 + LATENT_OFFSET, + mla_kv_half_append, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.attention.sparse_mla_flash_decode import ( # noqa: E501 + GLM_SOFTMAX_SCALE, + HEAD_DIM_QK, + KV_LORA_RANK, + MAX_TOP_K, + NUM_HEADS, + V_HEAD_DIM, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.attention.sparse_mla_flash_decode import ( # noqa: E501 + sparse_mla_flash_decode as op, +) +from tensorrt_llm._torch.attention.backends.interface import RopeParams +from tensorrt_llm._torch.attention.backends.sparse.dsa.cache_manager import DSACacheManagerV2 +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig, KvCacheConfig +from tensorrt_llm.mapping import Mapping + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or get_sm_version() != 107, + reason="the GLM-5.2 tp16 target's ops are certified on sm_107", +) + +#: log2(e) as an fp32 constant, as the KV append forms mla_bmm1_scale[1]. +_LOG2E = 1.4426950408889634 +#: Pool slots no sequence maps, and block-table columns past every sequence. +_SPARE_SLOTS = 5 +_SPARE_COLUMNS = 3 +#: Output rows past the query tokens, which the op must leave alone. +_SPARE_ROWS = 2 +#: Magnitudes of the random operands: KV rows ~ N(0, 1) as kv_a_layernorm leaves them, the +#: absorbed query ~ N(0, 0.25), v_b ~ N(0, 1 / 512). +_Q_STD = 0.5 +_VB_STD = 1.0 / math.sqrt(KV_LORA_RANK) +#: The relative L2 distance allowed between the reference and an attention whose +#: probabilities are not rounded to E4M3. E4M3 keeps 3 mantissa bits: a rounded probability is +#: off by up to 2**-4 relative (RMS 1.8-3.6 % across a binade), and on outputs that average +#: zero-mean rows the error stays at that relative size: the kernel measured 2.3 % against an +#: fp32 reference, and the bound is 4x that. A reference that mis-masks a slot, +#: drops a split or misses a scale is off by tens of percent. +_E4M3_P_REL_L2 = 0.1 + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """4 ulp of each token row's largest output, 1 ulp relative RMS, in BF16. + + The reference states the kernel's three roundings (P to E4M3 per 128-slot split, the merged + attention output to BF16, the v_b product to BF16), so what is left is fp32 accumulation + order: the QK and P.V sums inside mma.sync, the 16-way merge, exp2f, and the v_b sum. Those + move a value by about 1e-6 relative; that flips an E4M3 probability with probability about + 1e-5 per element (one flip in a cell of 64k probabilities at most), and one flipped + probability moves the outputs of its head by under 1 % of their magnitude (2 ulp of the + element, under 1 ulp of the row's largest); the final BF16 roundings flip by one ulp at + the same rate. A mirror with a wrong merge weight, a wrong split boundary or a dropped + head is off by many percent. + """ + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=1.0) + + +def _fill_e4m3_(dst: torch.Tensor, gen: torch.Generator, std: float = 1.0) -> None: + """Fill a contiguous uint8 view with E4M3 values of N(0, std^2) (finite: far below 448).""" + flat = dst.reshape(-1) + chunk = 1 << 26 + for start in range(0, flat.numel(), chunk): + stop = min(start + chunk, flat.numel()) + values = torch.randn(stop - start, generator=gen, device=flat.device) * std + flat[start:stop].copy_(values.to(torch.float8_e4m3fn).view(torch.uint8)) + + +def _select(kv_lens, seq_len_q: int, top_k: int, holes: bool, gen: torch.Generator) -> torch.Tensor: + """[T, top_k] int32 top-k positions as the DSA indexer lays them out. + + Row t of sequence s selects min(own + 1, top_k) positions out of 0 .. own (own = its + position): every token of this step it can see, the rest drawn from the earlier steps', + shuffled, then -1. With `holes`, about 1/64 of the selected slots become -1 in place. + """ + rows = [] + for kv_len in kv_lens: + first_new = kv_len - seq_len_q + for t in range(seq_len_q): + own = first_new + t + count = min(own + 1, top_k) + fresh = torch.arange(first_new, own + 1, dtype=torch.int64) + num_old = count - fresh.numel() + old = torch.randperm(first_new, generator=gen)[:num_old] + chosen = torch.cat((fresh, old)) + chosen = chosen[torch.randperm(chosen.numel(), generator=gen)] + if holes: + drop = torch.randperm(count, generator=gen)[: max(1, count // 64)] + chosen[drop] = -1 + row = torch.full((top_k,), -1, dtype=torch.int32) + row[:count] = chosen.to(torch.int32) + rows.append(row) + return torch.stack(rows) + + +def _global_rows( + positions: torch.Tensor, + table: torch.Tensor, + seq_len_q: int, + tokens_per_block: int, + page_index_scale: int, + layer_offset: int, +) -> torch.Tensor: + """The pool rows of `positions` with convert_req_index_to_global's arithmetic (-1 kept).""" + num_tokens = positions.shape[0] + seq = torch.arange(num_tokens, device=positions.device) // seq_len_q + position = positions.long() + block = position.clamp(min=0) // tokens_per_block + page = table[seq[:, None], block.clamp(max=table.shape[1] - 1)].long() + rows = ( + page * page_index_scale * tokens_per_block + + layer_offset * tokens_per_block + + position.clamp(min=0) % tokens_per_block + ) + return torch.where((position >= 0) & (page >= 0), rows, -1).to(torch.int32) + + +class _Case(NamedTuple): + kwargs: dict + num_tokens: int + global_rows: torch.Tensor + + +def _build(spec: dict, seed: int) -> _Case: + """Turn a cell spec into the op's arguments, over a pool of the spec's geometry.""" + gen_cpu = torch.Generator().manual_seed(seed) + gen = torch.Generator(device="cuda").manual_seed(seed) + tokens_per_block = spec["tokens_per_block"] + num_layers = spec["num_layers"] + layer_offset = spec["layer_offset"] + seq_len_q = spec["seq_len_q"] + top_k = spec["top_k"] + kv_lens = list(spec["kv_lens"]) + num_tokens = len(kv_lens) * seq_len_q + + blocks = [-(-kv_len // tokens_per_block) for kv_len in kv_lens] + num_slots = sum(blocks) + _SPARE_SLOTS + order = torch.randperm(num_slots, generator=gen_cpu).to(torch.int32) + table = torch.full((len(kv_lens), max(blocks) + _SPARE_COLUMNS), -1, dtype=torch.int32) + start = 0 + for s, num_blocks in enumerate(blocks): + table[s, :num_blocks] = order[start : start + num_blocks] + start += num_blocks + + pool = torch.empty( + num_slots, num_layers, tokens_per_block * HEAD_DIM_QK, dtype=torch.uint8, device="cuda" + ) + _fill_e4m3_(pool, gen) + positions = _select(kv_lens, seq_len_q, top_k, spec["holes"], gen_cpu) + global_rows = _global_rows( + positions, table, seq_len_q, tokens_per_block, num_layers, layer_offset + ) + + scale = torch.tensor(GLM_SOFTMAX_SCALE, dtype=torch.float32) + kwargs = dict( + q=( + torch.randn(num_tokens, NUM_HEADS, HEAD_DIM_QK, generator=gen, device="cuda") * _Q_STD + ).to(torch.float8_e4m3fn), + kv_pool=pool.view(torch.float8_e4m3fn), + indices=(positions if spec["remap"] else global_rows).cuda(), + seq_lens_kv=torch.tensor(kv_lens, dtype=torch.int32, device="cuda"), + mla_bmm1_scale=torch.stack((scale, scale * torch.tensor(_LOG2E))).cuda(), + mla_bmm2_scale=torch.ones(1, dtype=torch.float32, device="cuda"), + v_b_proj=( + torch.randn(NUM_HEADS, V_HEAD_DIM, KV_LORA_RANK, generator=gen, device="cuda") * _VB_STD + ).to(torch.bfloat16), + output=torch.full( + (num_tokens + _SPARE_ROWS, NUM_HEADS * V_HEAD_DIM), + float("nan"), + dtype=torch.bfloat16, + device="cuda", + ), + seq_len_q=seq_len_q, + top_k=top_k, + ) + if spec["remap"]: + kwargs.update( + block_table=table.cuda(), + tokens_per_block=tokens_per_block, + page_index_scale=num_layers, + layer_offset=layer_offset, + gather_before_wait=spec["gather_before_wait"], + ) + return _Case(kwargs=kwargs, num_tokens=num_tokens, global_rows=global_rows.cuda()) + + +def _unrounded_attention(case: _Case) -> torch.Tensor: + """[T, 1024] fp32: softmax attention over the selected rows without any intermediate + rounding, times v_b. An independent statement of what the op approximates.""" + kw = case.kwargs + num_tokens, seq_len_q, top_k = case.num_tokens, kw["seq_len_q"], kw["top_k"] + rows = case.global_rows.long() + token = torch.arange(num_tokens, device="cuda") + kv_len = kw["seq_lens_kv"][token // seq_len_q].long() + num_valid = torch.clamp(kv_len - (seq_len_q - 1 - token % seq_len_q), max=top_k) + valid = (rows >= 0) & (torch.arange(top_k, device="cuda")[None, :] < num_valid[:, None]) + pool_rows = kw["kv_pool"].view(torch.uint8).reshape(-1, HEAD_DIM_QK) + keys = pool_rows[rows.clamp(min=0)].view(torch.float8_e4m3fn).float().double() + queries = kw["q"].float().double() + scores = torch.einsum("thd,tjd->thj", queries, keys) * GLM_SOFTMAX_SCALE + scores = scores.masked_fill(~valid[:, None, :], float("-inf")) + weights = torch.nan_to_num(torch.softmax(scores, dim=-1), nan=0.0) + attn = torch.einsum("thj,tjd->thd", weights, keys[..., :KV_LORA_RANK]) + out = torch.einsum("thd,hvd->thv", attn, kw["v_b_proj"].double()) + return out.reshape(num_tokens, NUM_HEADS * V_HEAD_DIM).float() + + +def _seed(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by name + certifying(op) + case = _build(cell.spec, _seed(cell.why)) + output = case.kwargs["output"] + op(**case.kwargs) + torch.cuda.synchronize() + ref = op.reference(**case.kwargs) + + # The entry's reference is attention: within the E4M3-probability band of an unrounded + # one. Checked before the kernel gate, so a mirror that drifted from attention fails as + # such rather than as a kernel mismatch. + unrounded = _unrounded_attention(case) + rel_l2 = ((ref.float() - unrounded).norm() / unrounded.norm()).item() + assert rel_l2 <= _E4M3_P_REL_L2, f"reference is {rel_l2:.3%} from unrounded attention" + + _compare(output[: case.num_tokens], ref) + untouched = output[case.num_tokens :].float() + assert torch.isnan(untouched).all(), "rows past the tokens were written" + + +def test_capture_and_replay_match_the_eager_call() -> None: + """The shared-layer call (early gather) captured in a CUDA graph replays bit for bit, and + two eager calls agree bit for bit.""" + cell = op.CELLS[1] + assert cell.spec["gather_before_wait"] + case = _build(cell.spec, _seed(cell.why)) + output = case.kwargs["output"] + op(**case.kwargs) + torch.cuda.synchronize() + eager = output.clone() + output.fill_(float("nan")) + op(**case.kwargs) + torch.cuda.synchronize() + assert torch.equal(output[: case.num_tokens], eager[: case.num_tokens]), "not deterministic" + + graph = torch.cuda.CUDAGraph() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + with torch.cuda.graph(graph, stream=stream): + op(**case.kwargs) + torch.cuda.current_stream().wait_stream(stream) + for _ in range(3): + output.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(output[: case.num_tokens], eager[: case.num_tokens]) + + +def _guard_case(cell_index: int) -> _Case: + cell = op.CELLS[cell_index] + return _build(cell.spec, _seed(cell.why)) + + +def test_guard_refuses_a_selection_past_the_valid_count() -> None: + """A selected token behind the row's valid count is never read.""" + case = _guard_case(4) # short context: rows hold fewer valid slots than the top-k + seq_len_q = case.kwargs["seq_len_q"] + first_unread = int(case.kwargs["seq_lens_kv"][0]) - (seq_len_q - 1) # row 0's valid count + indices = case.kwargs["indices"] + assert first_unread < indices.shape[1] + indices[0, first_unread] = 0 + with pytest.raises(AssertionError, match="past the row's valid count"): + with validating(op): + op(**case.kwargs) + + +def test_guard_refuses_a_position_past_the_token_itself() -> None: + """A position after the query token's own reads a row it may not see.""" + case = _guard_case(0) + kv_len = int(case.kwargs["seq_lens_kv"][0]) + case.kwargs["indices"][0, 0] = kv_len - 1 # row 0 sits at kv_len - 8 + with pytest.raises(AssertionError, match="past the query token's own"): + with validating(op): + op(**case.kwargs) + + +def test_guard_refuses_a_table_narrower_than_the_kv() -> None: + """A table cut short masks live tokens instead of failing.""" + case = _guard_case(0) + kv_len = int(case.kwargs["seq_lens_kv"][0]) + needed = -(-kv_len // case.kwargs["tokens_per_block"]) + case.kwargs["block_table"] = case.kwargs["block_table"][:, : needed - 1] + with pytest.raises(AssertionError, match="columns; a sequence needs"): + with validating(op): + op(**case.kwargs) + + +def test_guard_refuses_a_hole_in_the_table() -> None: + """A -1 inside a sequence's blocks masks its tokens instead of failing.""" + case = _guard_case(0) + case.kwargs["block_table"][0, 3] = -1 + with pytest.raises(AssertionError, match="-1 inside"): + with validating(op): + op(**case.kwargs) + + +def test_guard_refuses_a_layer_offset_outside_the_slot() -> None: + """An offset at the page scale reads the next slot's first layer.""" + case = _guard_case(0) + case.kwargs["layer_offset"] = case.kwargs["page_index_scale"] + with pytest.raises(AssertionError, match="below page_index_scale"): + with validating(op): + op(**case.kwargs) + + +def test_guard_refuses_a_global_row_past_the_pool() -> None: + """Global rows are read as given.""" + case = _guard_case(3) # draft steps 1-6: global rows + assert "block_table" not in case.kwargs + num_rows = case.kwargs["kv_pool"].numel() // HEAD_DIM_QK + case.kwargs["indices"][0, 0] = num_rows + with pytest.raises(AssertionError, match="past the pool"): + with validating(op): + op(**case.kwargs) + + +class _DsaEnv: + """A DSACacheManagerV2 over GLM-5.2's latent rows (FP8, 576 per token, 64-token pages).""" + + TOKENS_PER_BLOCK = 64 + + def __init__(self, num_layers: int, max_batch_size: int, max_seq_len: int) -> None: + self.max_batch_size = max_batch_size + self.manager = DSACacheManagerV2( + kv_cache_config=KvCacheConfig( + max_tokens=max_batch_size * (max_seq_len + self.TOKENS_PER_BLOCK), + enable_block_reuse=False, + ), + kv_cache_type=CacheType.SELFKONLY, + num_layers=num_layers, + num_kv_heads=1, + head_dim=HEAD_DIM_QK, + tokens_per_block=self.TOKENS_PER_BLOCK, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=DataType.FP8, + sparse_attention_config=DeepSeekSparseAttentionConfig( + index_head_dim=128, index_n_heads=32, index_topk=MAX_TOP_K + ), + ) + rope = RopeParams( + dim=HEAD_DIM_QK - KV_LORA_RANK, + theta=8_000_000.0, + max_positions=max_seq_len, + duplicate_data=True, + ) + _, self.rotary_cos_sin = rope.create_rope_const_params() + + def generation_step(self, cached: list, seq_len_q: int) -> TrtllmAttentionMetadata: + """Requests 0 .. n-1 with `cached` tokens each, grown by one step of `seq_len_q`.""" + request_ids = list(range(len(cached))) + self.manager.add_dummy_requests(request_ids, token_nums=cached) + for rid in request_ids: + cache = self.manager.kv_cache_map[rid] + assert cache.resize(cache.capacity + seq_len_q) + metadata = TrtllmAttentionMetadata( + max_num_requests=self.max_batch_size, + max_num_tokens=8192, + kv_cache_manager=self.manager, + ) + metadata.seq_lens = torch.tensor([seq_len_q] * len(cached), dtype=torch.int) + metadata.num_contexts = 0 + metadata.request_ids = request_ids + metadata.prompt_lens = cached + metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=cached) + metadata.prepare() + return metadata + + def block_table(self, kv_lens: list) -> torch.Tensor: + """The DSA metadata's generation block table: pool slots, -1 past each sequence.""" + slots = self.manager.get_pool_block_indices( + len(kv_lens), request_ids=list(range(len(kv_lens))), num_contexts=0 + ) + blocks = torch.tensor([-(-n // self.TOKENS_PER_BLOCK) for n in kv_lens]) + width = int(blocks.max()) + table = slots[:, :width].to(torch.int32).clone() + table[torch.arange(width)[None, :] >= blocks[:, None]] = -1 + return table.cuda() + + def shutdown(self) -> None: + self.manager.shutdown() + + +def _convert( + block_table: torch.Tensor, positions: torch.Tensor, seq_len_q: int, page_params, tpb: int +) -> torch.Tensor: + """Main's pool rows of `positions` ([T, k] int32): convert_req_index_to_global.""" + page_index_scale, layer_offset = page_params + num_tokens = positions.shape[0] + req_idx = (torch.arange(num_tokens, device="cuda") // seq_len_q).to(torch.int32) + return torch.ops.trtllm.convert_req_index_to_global( + req_idx, + block_table, + positions, + tpb, + positions.shape[1], + page_index_scale * tpb, + layer_offset, + ) + + +def test_reads_the_rows_the_kv_append_writes_in_the_dsa_pool() -> None: + """Both ops against main's DSA cache layout. + + On a DSACacheManagerV2 pool, glm_mla_kv_half_append appends this step's 8 rows of one layer + through the attention metadata's block offsets; they must land where main's + convert_req_index_to_global places those positions. The flash decode then remaps the step's + top-k positions in the kernel, with and without the early gather, and must read exactly the + rows convert_req_index_to_global names: its output is gated against the reference run on + those global rows. + """ + seq_len_q, layer = 8, 2 + cached = [5_000] + kv_lens = [c + seq_len_q for c in cached] + num_tokens = len(cached) * seq_len_q + gen_cpu = torch.Generator().manual_seed(1) + gen = torch.Generator(device="cuda").manual_seed(1) + env = _DsaEnv(num_layers=4, max_batch_size=2, max_seq_len=8192) + try: + metadata = env.generation_step(cached, seq_len_q) + tpb = env.TOKENS_PER_BLOCK + pool = env.manager.get_unique_primary_pool() + pool_rows = pool.view(torch.uint8).reshape(-1, HEAD_DIM_QK) + _fill_e4m3_(pool_rows, gen) + block_table = env.block_table(kv_lens) + page_params = env.manager.get_primary_pool_page_index_params(layer) + + fused_a = torch.randn(num_tokens, 2624, generator=gen, device="cuda").to(torch.bfloat16) + latent = fused_a[:, LATENT_OFFSET : LATENT_OFFSET + HEAD_DIM_QK] + kv_norm_weight = (1.0 + 0.1 * torch.randn(KV_LORA_RANK, generator=gen, device="cuda")).to( + torch.bfloat16 + ) + bmm1 = torch.full((2,), float("nan"), dtype=torch.float32, device="cuda") + bmm2 = torch.full((1,), float("nan"), dtype=torch.float32, device="cuda") + append_args = ( + latent, + kv_norm_weight, + 1e-5, + env.rotary_cos_sin, + metadata.kv_lens_cuda_runtime, + 0, + seq_len_q, + metadata.kv_cache_block_offsets, + env.manager.kv_cache_pool_pointers, + env.manager.kv_cache_pool_mapping, + layer, + tpb, + env.manager.max_seq_len, + bmm1, + bmm2, + GLM_SOFTMAX_SCALE, + True, + ) + with validating(mla_kv_half_append): + mla_kv_half_append(*append_args) + torch.cuda.synchronize() + want_rows, positions, want_bmm1, want_bmm2 = mla_kv_half_append.reference(*append_args) + fresh = _convert( + block_table, positions.to(torch.int32)[:, None], seq_len_q, page_params, tpb + ) + got_rows = pool_rows[fresh[:, 0].long()].view(torch.float8_e4m3fn) + torch.testing.assert_close(got_rows.float(), want_rows.float(), rtol=2**-3, atol=2**-9) + assert torch.equal(bmm1, want_bmm1) and torch.equal(bmm2, want_bmm2) + + indices = _select(kv_lens, seq_len_q, MAX_TOP_K, False, gen_cpu).cuda() + global_rows = _convert(block_table, indices, seq_len_q, page_params, tpb) + q = ( + torch.randn(num_tokens, NUM_HEADS, HEAD_DIM_QK, generator=gen, device="cuda") * _Q_STD + ).to(torch.float8_e4m3fn) + v_b_proj = ( + torch.randn(NUM_HEADS, V_HEAD_DIM, KV_LORA_RANK, generator=gen, device="cuda") * _VB_STD + ).to(torch.bfloat16) + output = torch.empty( + num_tokens, NUM_HEADS * V_HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + ref = op.reference( + q, + pool, + global_rows, + metadata.kv_lens_cuda_runtime, + bmm1, + bmm2, + v_b_proj, + output, + seq_len_q, + MAX_TOP_K, + ) + for gather_before_wait in (False, True): + output.fill_(float("nan")) + with validating(op): + op( + q, + pool, + indices, + metadata.kv_lens_cuda_runtime, + bmm1, + bmm2, + v_b_proj, + output, + seq_len_q, + MAX_TOP_K, + block_table, + tpb, + page_params[0], + page_params[1], + gather_before_wait, + ) + torch.cuda.synchronize() + _compare(output, ref) + finally: + env.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_allgather_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_allgather_op_matrix.py new file mode 100644 index 000000000000..754b92766b87 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_allgather_op_matrix.py @@ -0,0 +1,243 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for the GLM-5.2 target's mnnvl_allgather entry. + +Its own launcher, like ``_mnnvl_hop_op_matrix.py`` (see ``_mnnvl_ranks``): + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _mnnvl_allgather_op_matrix.py + +The gather moves bits, so every gate is exact. Beyond the cell: special values (-0.0, infinities, +NaN payloads, the all-ones word), the rotation of the flags and the refill of the previous +buffer with the empty word, several private workspaces side by side (the gather's and the draft +steps'), CUDA graph replays, ranks arriving late, and the guards. The collected entry point is +``test_modeling_v2_mnnvl_allgather_op_matrix.py``. +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from typing import Any, List, Tuple + +import _mnnvl_ranks as ranks +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "mnnvl_allgather requires CUDA devices" + +DEADLINE_S = 1200 +SHARD = 9680 +# One Lamport buffer of the logits gather: the gathered logits of up to 54 rows of the +# 154,880-entry BF16 vocabulary. +GATHER_BUFFER_BYTES = 16 * 1024 * 1024 +EMPTY_WORD = -1 + +# Bound in the rank body: the entry and the workspace helpers import tensorrt_llm. +entry: Any = None +ws: Any = None +validating: Any = None +RANK = 0 +WORLD = 1 +MAPPING: Any = None +WORKSPACE: Any = None + + +def _shard(rows: int, seed: int, rank: int, shard: int = SHARD) -> torch.Tensor: + return ranks.normal((rows, shard), seed * 1009 + rank, scale=4.0) + + +def _gather(x: torch.Tensor, workspace: Any = None) -> torch.Tensor: + workspace = WORKSPACE if workspace is None else workspace + return entry.mnnvl_allgather(x, workspace.comm_buffer, workspace.buffer_flags, WORLD) + + +def _check(out: torch.Tensor, per_rank: List[torch.Tensor], where: str) -> None: + ranks.assert_bitwise(out, entry.mnnvl_allgather.reference(per_rank), where) + + +def check_certified_cells() -> None: + """Every cell: the concatenation bit for bit, the same on every rank.""" + for i, cell in enumerate(entry.mnnvl_allgather.CELLS): + rows, shard = cell.spec["rows"], cell.spec["shard"] + seed = 100 + 13 * i + per_rank = [_shard(rows, seed, r, shard) for r in range(WORLD)] + with validating(entry.mnnvl_allgather): + out = entry.mnnvl_allgather( + per_rank[RANK], WORKSPACE.comm_buffer, WORKSPACE.buffer_flags, WORLD + ) + assert out.shape == (rows, WORLD * shard) and out.dtype == torch.float32, cell.why + _check(out, per_rank, cell.why) + ranks.assert_same_on_every_rank(out, cell.why) + dist.barrier() + + +def check_special_values_travel_unchanged() -> None: + """-0.0, infinities, NaN payloads and subnormals arrive bit for bit; a word of all ones (two + BF16 NaNs) arrives with its lowest bit cleared, still two NaNs.""" + specials = torch.tensor( + [0x8000, 0x0000, 0x7F80, 0xFF80, 0x7FC1, 0xFFA5, 0x0001, 0x8001, 0xFFFF, 0xFFFF], + dtype=torch.int32, + ).to(torch.int16) + per_rank = [] + for r in range(WORLD): + x = _shard(8, 2000, r) + x.view(torch.int16)[r % 8, : specials.numel()] = specials.cuda() + per_rank.append(x) + expected_inputs = [] + for x in per_rank: + words = x.clone().view(torch.int32) + words[words == EMPTY_WORD] = EMPTY_WORD - 1 + expected_inputs.append(words.view(torch.bfloat16)) + out = _gather(per_rank[RANK]) + ranks.assert_bitwise(out, entry.mnnvl_allgather.reference(expected_inputs), "special values") + + +def check_flags_rotate_and_the_previous_buffer_is_emptied() -> None: + """Each call rotates the flags; the next one fills the buffer it used with the empty word.""" + previous = None + for step, rows in enumerate((8, 1, 3, 8, 2, 8, 1)): + seed = 3000 + step + per_rank = [_shard(rows, seed, r) for r in range(WORLD)] + before = ranks.lamport_flags(WORKSPACE.buffer_flags) + _check(_gather(per_rank[RANK]), per_rank, f"step {step}") + after = ranks.lamport_flags(WORKSPACE.buffer_flags) + gathered_bytes = rows * WORLD * SHARD * torch.bfloat16.itemsize + ranks.assert_rotated(before, after, gathered_bytes, f"step {step}") + if previous is not None: + index, words = previous + emptied = WORKSPACE.comm_buffer.view(3, -1)[index, :words] + assert bool((emptied == EMPTY_WORD).all()), f"step {step}: buffer {index} not emptied" + previous = (before[0], gathered_bytes // 4) + dist.barrier() + + +def check_private_workspaces_side_by_side() -> None: + """The draft steps' exchange workspaces live beside the gather's; calls alternate freely.""" + # The sizes of the draft_sample_glue workspace (8 rows of 16-byte packets per rank) and of the + # draft_ehproj_table one (the packet, a BF16 hidden row and the rank's FP32 hidden-half share). + glue = ws.LamportExchangeWorkspace(MAPPING, 8 * WORLD * 16) + table = ws.LamportExchangeWorkspace(MAPPING, 8 * WORLD * (16 + 2 * 6144 + 4 * 6144 // WORLD)) + shards = {id(glue): 8, id(table): 6144 // WORLD, id(WORKSPACE): SHARD} + for step in range(6): + workspace = (WORKSPACE, glue, table)[step % 3] + shard = shards[id(workspace)] + rows = 1 if workspace is glue else 8 + per_rank = [_shard(rows, 4000 + step, r, shard) for r in range(WORLD)] + _check(_gather(per_rank[RANK], workspace), per_rank, f"workspace {step % 3} step {step}") + dist.barrier() + + +def check_cuda_graph_replays() -> None: + """A captured gather reads its input buffer and the rotating flags at every replay.""" + rows = 8 + x_static = _shard(rows, 5000, RANK).clone() + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + _gather(x_static) + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = _gather(x_static) + for replay in range(7): + per_rank = [_shard(rows, 5100 + replay, r) for r in range(WORLD)] + x_static.copy_(per_rank[RANK]) + torch.cuda.synchronize() + dist.barrier() + graph.replay() + torch.cuda.synchronize() + _check(out, per_rank, f"replay {replay}") + dist.barrier() + del graph + torch.cuda.synchronize() + + +def check_late_ranks() -> None: + """Ranks reach back-to-back gathers at different times.""" + issued: List[Tuple[List[torch.Tensor], torch.Tensor]] = [] + for step in range(6): + rows = (8, 1, 5)[step % 3] + per_rank = [_shard(rows, 6000 + step, r) for r in range(WORLD)] + torch.cuda._sleep(((RANK * 5 + step) % WORLD) * 400_000) + issued.append((per_rank, _gather(per_rank[RANK]))) + torch.cuda.synchronize() + for step, (per_rank, out) in enumerate(issued): + _check(out, per_rank, f"late ranks step {step}") + + +def check_the_op_refuses_rows_past_the_buffer() -> None: + """Gathered rows that exceed one Lamport buffer are refused before any rank stores them.""" + rows = GATHER_BUFFER_BYTES // (WORLD * SHARD * torch.bfloat16.itemsize) + 1 + x = torch.zeros((rows, SHARD), dtype=torch.bfloat16, device="cuda") + try: + _gather(x) + except RuntimeError as exc: + assert "exceed the Lamport buffer" in str(exc), exc + else: + raise AssertionError(f"the op took {rows} rows past the buffer") + + +def check_the_guard_refuses_the_all_reduce_workspace() -> None: + """The all-reduce workspace marks empty words with -0.0: is_valid refuses it.""" + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, 8, 6144) + x = _shard(8, 7000, RANK) + try: + with validating(entry.mnnvl_allgather): + entry.mnnvl_allgather(x, comm_buffer, buffer_flags, WORLD) + except AssertionError as exc: + assert "LamportExchangeWorkspace" in str(exc), exc + else: + raise AssertionError("is_valid admitted the all-reduce workspace") + + +def check_the_guard_refuses_flags_of_another_workspace() -> None: + foreign = WORKSPACE.buffer_flags.clone() + foreign.view(torch.int32)[2] += 16 + x = _shard(8, 7100, RANK) + try: + with validating(entry.mnnvl_allgather): + entry.mnnvl_allgather(x, WORKSPACE.comm_buffer, foreign, WORLD) + except AssertionError as exc: + assert "flags of comm_buffer's workspace" in str(exc), exc + else: + raise AssertionError("is_valid admitted foreign flags") + + +CHECKS = ( + check_certified_cells, + check_special_values_travel_unchanged, + check_flags_rotate_and_the_previous_buffer_is_emptied, + check_private_workspaces_side_by_side, + check_cuda_graph_replays, + check_late_ranks, + check_the_op_refuses_rows_past_the_buffer, + check_the_guard_refuses_the_all_reduce_workspace, + check_the_guard_refuses_flags_of_another_workspace, +) + + +def _run_one_rank() -> None: + global entry, ws, validating, RANK, WORLD, MAPPING, WORKSPACE + RANK, WORLD, MAPPING = ranks.init_rank() + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.comm import ( # noqa: E501 + _mnnvl_workspace, + mnnvl_allgather, + ) + + entry, ws = mnnvl_allgather, _mnnvl_workspace + WORKSPACE = ws.LamportExchangeWorkspace(MAPPING, GATHER_BUFFER_BYTES) + # Run as a script, not collected: the shared _validating module is reached the plain way. + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating as _validating_cm + + validating = _validating_cm + ranks.run_checks(CHECKS, RANK) + + +if __name__ == "__main__": + if ranks.WORKER_FLAG in sys.argv: + _run_one_rank() + ranks.spawn_ranks(Path(__file__).resolve(), DEADLINE_S) + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_hop_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_hop_op_matrix.py new file mode 100644 index 000000000000..f8847b2712f7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_hop_op_matrix.py @@ -0,0 +1,420 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for the GLM-5.2 target's mnnvl_hop entry. + +A collective cannot be exercised in one process, so this script is its own launcher (see +``_mnnvl_ranks``): run plainly it spawns one rank per device in CUDA_VISIBLE_DEVICES, + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _mnnvl_hop_op_matrix.py + +and ``--rank-worker`` runs one rank of a world started elsewhere (the 16-rank c1 world). The +checks are one fixed sequence on one MNNVL all-reduce workspace, which each check leaves rotated +for the next; collected as tests they would run at world size 1, so neither this file's name nor +its functions match pytest's patterns. The collected entry point is +``test_modeling_v2_mnnvl_hop_op_matrix.py`` beside this file. + +Beyond the entry's cells: exact sums, the Lamport flags, -0.0 rows, the published shard, +interleaving with the built-in MNNVL all-reduce on its workspace, CUDA graph replays, ranks +arriving late, and the guards. On a world whose one-shot all-reduce stays under 1 MiB the built-in +op computes the same values in the same order, so the hop matches it bit for bit. +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from typing import Any, List, Tuple + +import _mnnvl_ranks as ranks +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "mnnvl_hop requires CUDA devices" + +DEADLINE_S = 1200 +HIDDEN = 6144 +EPS = 1e-5 + +# Bound in the rank body: the entry and the workspace helpers import tensorrt_llm. +entry: Any = None +ws: Any = None +validating: Any = None +fusion_op: Any = None +RANK = 0 +WORLD = 1 +MAPPING: Any = None + + +def _rows(rows: int, seed: int, rank: int) -> torch.Tensor: + """One rank's partial rows: their sum over the ranks has unit variance.""" + return ranks.normal((rows, HIDDEN), seed * 1009 + rank, scale=WORLD**-0.5) + + +def _residual(rows: int, seed: int, zero: bool = False) -> torch.Tensor: + """The residual stream: the same on every rank, as on the target.""" + if zero: + return torch.zeros((rows, HIDDEN), dtype=torch.bfloat16, device="cuda") + return ranks.normal((rows, HIDDEN), seed * 1009 + 500) + + +def _gamma(seed: int) -> torch.Tensor: + """An RMSNorm weight near 1, the same on every rank.""" + return (1.0 + ranks.normal(HIDDEN, seed, scale=0.1, dtype=torch.float32)).to(torch.bfloat16) + + +def _compare(norm, residual_out, ref_norm, ref_residual, where: str) -> None: + """Bands of the hop against the FP32 reference. + + residual_out: the kernel rounds the FP32 sum over the ranks to BF16 and adds the residual in + BF16, two roundings of half an ulp, each of a value at most the row's largest (inputs scaled + so that the sum and the residual are of one size): at most about 2 ulp of the row scale per + element, 3 allowed; about 0.6 ulp of relative RMS, 1.5 allowed. + + norm_out: residual_out's error carried through the scale (gamma within 0.6-1.4), the sum of + squares of the BF16 values (each square rounded to BF16, below 1/4 ulp on the scale) and the + output rounding (half an ulp): at most about 3.5 ulp per element, 5 allowed; about 0.7 ulp of + relative RMS, 1.5 allowed. Derived, not fitted: a misplaced slot or rank misses by orders. + """ + from _validating import assert_within_ulp + + try: + assert_within_ulp(residual_out, ref_residual, element_ulp=3.0, rms_ulp=1.5) + assert_within_ulp(norm, ref_norm, element_ulp=5.0, rms_ulp=1.5) + except AssertionError as exc: + raise AssertionError(f"{where}: {exc}") from exc + + +def _builtin_one_shot(rows: int) -> bool: + """Whether the built-in MNNVL all-reduce runs one-shot here, in the hops' reduction order.""" + return rows * HIDDEN * WORLD * torch.bfloat16.itemsize <= ranks.BUILTIN_ONE_SHOT_BYTES + + +def _builtin_hop(x, residual, gamma, comm_buffer, buffer_flags) -> Tuple[torch.Tensor, ...]: + """The built-in MNNVL residual + RMSNorm all-reduce on the same workspace.""" + norm, residual_out = torch.ops.trtllm.mnnvl_fusion_allreduce( + x, gamma, residual, EPS, comm_buffer, buffer_flags, True, None, fusion_op + ) + return norm, residual_out + + +def _word_offsets(rows: int) -> torch.Tensor: + """Offsets, in 32-bit words of a Lamport buffer, of this rank's slot of every row.""" + tokens = torch.arange(rows, device="cuda").view(-1, 1) + pairs = torch.arange(0, HIDDEN, 2, device="cuda").view(1, -1) + return (((tokens * WORLD + RANK) * HIDDEN + pairs) // 2).view(-1) + + +def _publish(x: torch.Tensor, offsets: torch.Tensor | None = None) -> None: + """Store this rank's rows into every rank's current Lamport buffer, as the o_proj publish does. + + The GEMM's epilogue reads the current buffer from the flags and stores its BF16 output, -0.0 + as +0.0, at [token][rank][column]. Done here with 32-bit stores through the multicast view: + a word is the unit the hop's poll tests, so a half-written slot is never taken for a full one. + The index comes from the flags on the device, so this also runs inside a CUDA graph. + """ + rows = x.shape[0] + lamport, flags = ws.publish_target(MAPPING, rows, HIDDEN) + words = lamport.view(torch.int32) + if offsets is None: + offsets = _word_offsets(rows) + base = flags[0].long() * (flags[2].long() // 4) + payload = torch.where(x == 0, torch.zeros_like(x), x).contiguous().view(torch.int32) + words.index_copy_(0, base + offsets, payload.view(-1)) + + +def _hop(x, residual, gamma, shard_published: bool) -> Tuple[torch.Tensor, torch.Tensor]: + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, x.shape[0], HIDDEN) + if shard_published: + _publish(x) + return entry.mnnvl_hop(x, residual, gamma, EPS, comm_buffer, buffer_flags, shard_published) + + +def check_certified_cells() -> None: + """Every cell: the band against the reference, the same bits on every rank, and the bits of + the built-in one-shot all-reduce where it runs one-shot.""" + gamma = _gamma(7) + for i, cell in enumerate(entry.mnnvl_hop.CELLS): + spec = cell.spec + rows, published, zero = spec["rows"], spec["shard_published"], spec["zero_residual"] + assert spec["hidden"] == HIDDEN, cell.why + seed = 100 + 13 * i + per_rank = [_rows(rows, seed, r) for r in range(WORLD)] + x, residual = per_rank[RANK], _residual(rows, seed, zero) + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, rows, HIDDEN) + if published: + _publish(x) + with validating(entry.mnnvl_hop): + norm, residual_out = entry.mnnvl_hop( + x, residual, gamma, EPS, comm_buffer, buffer_flags, published + ) + ref_norm, ref_residual = entry.mnnvl_hop.reference(per_rank, residual, gamma, EPS) + _compare(norm, residual_out, ref_norm, ref_residual, cell.why) + ranks.assert_same_on_every_rank(norm, f"{cell.why}: norm_out") + ranks.assert_same_on_every_rank(residual_out, f"{cell.why}: residual_out") + if _builtin_one_shot(rows): + builtin_norm, builtin_residual = _builtin_hop( + x, residual, gamma, comm_buffer, buffer_flags + ) + ranks.assert_bitwise(norm, builtin_norm, f"{cell.why}: norm_out vs built-in") + ranks.assert_bitwise(residual_out, builtin_residual, f"{cell.why}: residual_out") + dist.barrier() + + +def check_exact_sums_are_exact() -> None: + """Rows on the 1/8 grid sum exactly: residual_out equals the reference bit for bit. + + Any slot read from the wrong rank, row or column, or summed twice, moves an exact result. + """ + gamma = _gamma(11) + for rows, published in ((8, False), (8, True), (1, True), (3, False)): + seed = 2000 + rows + 7 * published + # |row| <= 1 per rank, |sum| <= 16 on 16 ranks, |residual| <= 8: all exact in BF16. + per_rank = [ranks.eighths((rows, HIDDEN), seed * 31 + r, bound=8) for r in range(WORLD)] + residual = ranks.eighths((rows, HIDDEN), seed * 31 + 97, bound=64) + norm, residual_out = _hop(per_rank[RANK], residual, gamma, published) + ref_norm, ref_residual = entry.mnnvl_hop.reference(per_rank, residual, gamma, EPS) + where = f"exact rows={rows} published={published}" + ranks.assert_bitwise(residual_out, ref_residual, where) + _compare(norm, residual_out, ref_norm, ref_residual, where) + dist.barrier() + + +def check_lamport_flags_rotate() -> None: + """After each hop the current buffer is the dirty one, with the hop's bytes to clear.""" + gamma = _gamma(13) + for step in range(4): + rows = (8, 1, 5, 8)[step] + published = step % 2 == 1 + _, buffer_flags = ws.hop_workspace(MAPPING, rows, HIDDEN) + before = ranks.lamport_flags(buffer_flags) + _hop(_rows(rows, 3000 + step, RANK), _residual(rows, 3000 + step), gamma, published) + after = ranks.lamport_flags(buffer_flags) + bytes_to_clear = rows * HIDDEN * WORLD * torch.bfloat16.itemsize + ranks.assert_rotated(before, after, bytes_to_clear, f"step {step}") + dist.barrier() + + +def check_negative_zero_rows_arrive_as_positive_zero() -> None: + """-0.0 marks an empty slot; rows of -0.0 must neither hang the exchange nor sum to -0.0.""" + gamma = _gamma(17) + for published in (False, True): + x = torch.full((8, HIDDEN), -0.0, dtype=torch.bfloat16, device="cuda") + residual = ranks.eighths((8, HIDDEN), 4000, bound=64) + # A residual of -0.0 plus a sum of +0.0 is +0.0, and any other residual is itself. + residual = torch.where(residual == 0, torch.ones_like(residual), residual) + _, residual_out = _hop(x, residual, gamma, published) + ranks.assert_bitwise(residual_out, residual, f"-0.0 rows published={published}") + dist.barrier() + + +def check_interleaved_with_the_builtin_all_reduce() -> None: + """The hops and the built-in MNNVL all-reduce (one-shot and two-shot) share one workspace. + + They read, clear and rotate the same flags; every result stays right whichever kernel ran + before it. + """ + gamma = _gamma(19) + for round_ in range(3): + seed = 5000 + 100 * round_ + for kind, rows in ( + ("hop", 8), + ("builtin", 2), + ("published", 1), + ("builtin_two_shot", 64), + ("published", 8), + ("hop", 1), + ): + per_rank = [_rows(rows, seed + rows, r) for r in range(WORLD)] + residual = _residual(rows, seed + rows) + if kind.startswith("builtin"): + # 64 rows are two-shot for the built-in all-reduce; a workspace that holds the + # rows one-shot holds them two-shot too. + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, rows, HIDDEN) + norm, residual_out = _builtin_hop( + per_rank[RANK], residual, gamma, comm_buffer, buffer_flags + ) + else: + norm, residual_out = _hop(per_rank[RANK], residual, gamma, kind == "published") + ref_norm, ref_residual = entry.mnnvl_hop.reference(per_rank, residual, gamma, EPS) + _compare(norm, residual_out, ref_norm, ref_residual, f"round {round_} {kind} {rows}") + dist.barrier() + + +def check_cuda_graph_replays() -> None: + """A captured decode-step chain replays right: publish, published hop, hop, built-in hop. + + The flags rotate on the device, so every replay runs on the buffers the previous kernel left + current; the published rows land where the flags say at replay time. + """ + gamma = _gamma(23) + rows = 8 + x_static = _rows(rows, 6000, RANK).clone() + residual_static = _residual(rows, 6000).clone() + x2_static = _rows(1, 6001, RANK).clone() + offsets = _word_offsets(rows) + + def chain(): + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, rows, HIDDEN) + _publish(x_static, offsets) + norm1, res1 = entry.mnnvl_hop( + x_static, residual_static, gamma, EPS, comm_buffer, buffer_flags, True + ) + norm2, res2 = entry.mnnvl_hop(norm1, res1, gamma, EPS, comm_buffer, buffer_flags, False) + comm_buffer1, buffer_flags1 = ws.hop_workspace(MAPPING, 1, HIDDEN) + norm3, res3 = entry.mnnvl_hop( + x2_static, res2[:1].contiguous(), gamma, EPS, comm_buffer1, buffer_flags1, False + ) + return norm1, res1, norm2, res2, norm3, res3 + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + chain() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs = chain() + for replay in range(7): + seed = 6100 + 10 * replay + per_rank = [_rows(rows, seed, r) for r in range(WORLD)] + per_rank1 = [_rows(1, seed + 1, r) for r in range(WORLD)] + x_static.copy_(per_rank[RANK]) + residual_static.copy_(_residual(rows, seed)) + x2_static.copy_(per_rank1[RANK]) + torch.cuda.synchronize() + dist.barrier() + graph.replay() + torch.cuda.synchronize() + norm1, res1, norm2, res2, norm3, res3 = outputs + ref = entry.mnnvl_hop.reference(per_rank, residual_static, gamma, EPS) + _compare(norm1, res1, *ref, f"replay {replay} published hop") + # The second hop sums norm1 over the ranks: build its inputs from every rank's norm1. + gathered = [torch.empty_like(norm1) for _ in range(WORLD)] + dist.all_gather(gathered, norm1) + ref = entry.mnnvl_hop.reference(gathered, res1, gamma, EPS) + _compare(norm2, res2, *ref, f"replay {replay} hop") + ref = entry.mnnvl_hop.reference(per_rank1, res2[:1], gamma, EPS) + _compare(norm3, res3, *ref, f"replay {replay} one-row hop") + ranks.assert_same_on_every_rank(norm3, f"replay {replay} one-row hop") + dist.barrier() + del graph + torch.cuda.synchronize() + + +def check_late_ranks() -> None: + """Ranks reach the hops at different times; back-to-back hops still exchange the right rows. + + Rank r sleeps on its stream before each hop, so the early ranks poll while the late ones + have not published, and a fast rank's next hop publishes into a buffer a slow rank has only + just emptied. + """ + gamma = _gamma(29) + issued: List[Tuple[Any, ...]] = [] + for step in range(6): + rows = (8, 1)[step % 2] + published = step % 3 == 0 + seed = 7000 + step + per_rank = [_rows(rows, seed, r) for r in range(WORLD)] + residual = _residual(rows, seed) + torch.cuda._sleep(((RANK * 7 + step * 3) % WORLD) * 400_000) + issued.append((per_rank, residual, *_hop(per_rank[RANK], residual, gamma, published))) + torch.cuda.synchronize() + for step, (per_rank, residual, norm, residual_out) in enumerate(issued): + ref_norm, ref_residual = entry.mnnvl_hop.reference(per_rank, residual, gamma, EPS) + _compare(norm, residual_out, ref_norm, ref_residual, f"late ranks step {step}") + + +def check_the_op_refuses_unsupported_shapes() -> None: + """Shapes outside hop_supported are refused loudly, before any rank enters the exchange.""" + assert ws.hop_supported(8, HIDDEN, WORLD) + assert ws.hop_supported(1, HIDDEN, WORLD) + # On 16 ranks: 8 rows of 6144 are the 1.5 MiB one-shot limit; 9 rows are above it. + assert ws.hop_supported(8, HIDDEN, 16) and not ws.hop_supported(9, HIDDEN, 16) + # 6152 columns are 769 eight-element vectors: no cluster of 2 or more CTAs splits them. + assert not ws.hop_supported(1, HIDDEN + 8, WORLD) + assert not ws.hop_supported(1, HIDDEN, 32) + too_many = 1 + while ws.hop_supported(too_many, HIDDEN, WORLD): + too_many += 1 + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, 1, HIDDEN) + for rows, hidden in ((too_many, HIDDEN), (1, HIDDEN + 8)): + x = torch.zeros((rows, hidden), dtype=torch.bfloat16, device="cuda") + g = torch.ones(hidden, dtype=torch.bfloat16, device="cuda") + try: + entry.mnnvl_hop(x, x, g, EPS, comm_buffer, buffer_flags, False) + except RuntimeError as exc: + assert "unsupported shape" in str(exc), exc + else: + raise AssertionError(f"the op took {rows} x {hidden} rows") + + +def check_the_guard_refuses_flags_of_another_workspace() -> None: + """Flags that do not describe comm_buffer's buffers are refused by is_valid.""" + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, 8, HIDDEN) + foreign = buffer_flags.clone() + foreign.view(torch.int32)[2] += 16 + x = _rows(8, 8000, RANK) + try: + with validating(entry.mnnvl_hop): + entry.mnnvl_hop(x, x, _gamma(37), EPS, comm_buffer, foreign, False) + except AssertionError as exc: + assert "flags of comm_buffer's workspace" in str(exc), exc + else: + raise AssertionError("is_valid admitted foreign flags") + + +def check_the_guard_refuses_a_workspace_of_another_empty_value() -> None: + """A workspace view that is not BF16 (an exchange workspace's words) is refused.""" + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, 8, HIDDEN) + x = _rows(8, 8100, RANK) + try: + with validating(entry.mnnvl_hop): + entry.mnnvl_hop( + x, x, _gamma(41), EPS, comm_buffer.view(torch.int32), buffer_flags, False + ) + except AssertionError as exc: + assert "BF16 view" in str(exc), exc + else: + raise AssertionError("is_valid admitted a non-BF16 workspace view") + + +CHECKS = ( + check_certified_cells, + check_exact_sums_are_exact, + check_lamport_flags_rotate, + check_negative_zero_rows_arrive_as_positive_zero, + check_interleaved_with_the_builtin_all_reduce, + check_cuda_graph_replays, + check_late_ranks, + check_the_op_refuses_unsupported_shapes, + check_the_guard_refuses_flags_of_another_workspace, + check_the_guard_refuses_a_workspace_of_another_empty_value, +) + + +def _run_one_rank() -> None: + global entry, ws, validating, fusion_op, RANK, WORLD, MAPPING + RANK, WORLD, MAPPING = ranks.init_rank() + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.comm import ( # noqa: E501 + _mnnvl_workspace, + mnnvl_hop, + ) + from tensorrt_llm.functional import AllReduceFusionOp + + entry, ws = mnnvl_hop, _mnnvl_workspace + fusion_op = int(AllReduceFusionOp.RESIDUAL_RMS_NORM) + # Run as a script, not collected: the shared _validating module is reached the plain way. + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating as _validating_cm + + validating = _validating_cm + ranks.run_checks(CHECKS, RANK) + + +if __name__ == "__main__": + if ranks.WORKER_FLAG in sys.argv: + _run_one_rank() + ranks.spawn_ranks(Path(__file__).resolve(), DEADLINE_S) + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_moe_finalize_hop_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_moe_finalize_hop_op_matrix.py new file mode 100644 index 000000000000..a6c3e5478206 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_moe_finalize_hop_op_matrix.py @@ -0,0 +1,353 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for the GLM-5.2 target's mnnvl_moe_finalize_hop entry. + +Its own launcher, like ``_mnnvl_hop_op_matrix.py`` (see ``_mnnvl_ranks``): + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _mnnvl_moe_finalize_hop_op_matrix.py + +The fused kernel must equal the two steps it replaces bit for bit: the TRTLLM-Gen finalize with +the shared experts' output as addend (computed here exactly in torch: a BF16 weight times a BF16 +row is exact in FP32, so the separate multiply and add are the FMA), followed by ``mnnvl_hop`` on +its rows. Beyond the cells: empty (-1) slots, FC2 rows wider than the hidden size, interleaving with +the plain hops on one workspace, CUDA graph replays, and the index guard. The collected entry +point is ``test_modeling_v2_mnnvl_moe_finalize_hop_op_matrix.py``. +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from typing import Any, List, Tuple + +import _mnnvl_ranks as ranks +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "mnnvl_moe_finalize_hop requires CUDA devices" + +DEADLINE_S = 1200 +HIDDEN = 6144 +EPS = 1e-5 +TOP_K = 8 + +# Bound in the rank body: the entries and the workspace helpers import tensorrt_llm. +entry: Any = None +hop: Any = None +ws: Any = None +validating: Any = None +RANK = 0 +WORLD = 1 +MAPPING: Any = None + +MoeInputs = Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] + + +def _moe_inputs( + rows: int, seed: int, rank: int, routing: str, pitch: int = HIDDEN, empty_slots: bool = False +) -> MoeInputs: + """One rank's ``(gemm2_output, expert_weights, expanded_idx_to_permuted_idx, addend)``. + + ``top8``: rows * 8 FC2 rows plus 5 padding rows, a random permutation of them over the + (token, k) slots, positive routing weights; with ``empty_slots`` some slots are -1. + ``routed_row``: the MTP layer's add in the hop, one FC2 row per token (its routed output) in + slot 0 with weight 1, slots 1-7 empty. Scaled so that the finalized rows summed over the ranks + have unit variance. + """ + scale = WORLD**-0.5 + base = seed * 7919 + rank * 131 + addend = ranks.normal((rows, HIDDEN), base + 1, scale=0.5 * scale) + if routing == "routed_row": + gemm2_output = ranks.normal((rows, pitch), base + 2, scale=0.5 * scale) + weights = torch.zeros((rows, TOP_K), dtype=torch.bfloat16, device="cuda") + weights[:, 0] = 1.0 + index = torch.full((rows, TOP_K), -1, dtype=torch.int32, device="cuda") + index[:, 0] = torch.arange(rows, dtype=torch.int32, device="cuda") + return gemm2_output, weights, index, addend + assert routing == "top8", routing + fc2_rows = rows * TOP_K + 5 + gemm2_output = ranks.normal((fc2_rows, pitch), base + 3, scale=scale) + gen = torch.Generator(device="cuda").manual_seed(base + 4) + index = torch.randperm(fc2_rows, generator=gen, device="cuda")[: rows * TOP_K] + index = index.to(torch.int32).view(rows, TOP_K).contiguous() + if empty_slots: + drop = torch.rand((rows, TOP_K), generator=gen, device="cuda") < 0.3 + index = torch.where(drop, torch.full_like(index, -1), index) + weights = (torch.rand((rows, TOP_K), generator=gen, device="cuda") * 0.25).to(torch.bfloat16) + return gemm2_output, weights, index, addend + + +def _finalized_rows(inputs: MoeInputs) -> torch.Tensor: + """The TRTLLM-Gen finalize with an addend, bit for bit: the routed sum in FP32 in k order + (each weight times row exact in FP32, so the add is the FMA's rounding), rounded to BF16, the + addend added in FP32 and rounded once more.""" + gemm2_output, weights, index, addend = inputs + routed = torch.zeros(addend.shape, dtype=torch.float32, device="cuda") + for k in range(TOP_K): + slot = index[:, k] + taken = (slot >= 0).view(-1, 1) + rows = gemm2_output[slot.clamp(min=0).long(), :HIDDEN].float() + product = weights[:, k].float().view(-1, 1) * rows + routed = torch.where(taken, routed + product, routed) + return (addend.float() + routed.to(torch.bfloat16).float()).to(torch.bfloat16) + + +def _residual(rows: int, seed: int) -> torch.Tensor: + """The residual stream: the same on every rank, as on the target.""" + return ranks.normal((rows, HIDDEN), seed * 1009 + 500) + + +def _gamma(seed: int) -> torch.Tensor: + return (1.0 + ranks.normal(HIDDEN, seed, scale=0.1, dtype=torch.float32)).to(torch.bfloat16) + + +def _compare(norm, residual_out, ref_norm, ref_residual, where: str) -> None: + """Bands of the fused hop against the FP32 reference. + + On top of mnnvl_hop's two roundings (the sum, the residual add) each rank rounds its routed + sum and then the addend's sum to BF16 before the exchange: two more half-ulp roundings of a + rank's row, whose sum over the ranks (independent signs) stays near one ulp of the summed + row's scale. residual_out: 4 ulp per element and 2 ulp of relative RMS allowed; norm_out: 6 + and 2. Derived, not fitted; the fused-vs-two-step gate in the checks is exact. + """ + from _validating import assert_within_ulp + + try: + assert_within_ulp(residual_out, ref_residual, element_ulp=4.0, rms_ulp=2.0) + assert_within_ulp(norm, ref_norm, element_ulp=6.0, rms_ulp=2.0) + except AssertionError as exc: + raise AssertionError(f"{where}: {exc}") from exc + + +def _fused(inputs: MoeInputs, residual, gamma) -> Tuple[torch.Tensor, torch.Tensor]: + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, residual.shape[0], HIDDEN) + gemm2_output, weights, index, addend = inputs + return entry.mnnvl_moe_finalize_hop( + gemm2_output, weights, index, addend, residual, gamma, EPS, comm_buffer, buffer_flags + ) + + +def _two_step(inputs: MoeInputs, residual, gamma) -> Tuple[torch.Tensor, torch.Tensor]: + """The finalize, then the plain hop on its rows.""" + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, residual.shape[0], HIDDEN) + rows = _finalized_rows(inputs) + return hop.mnnvl_hop(rows, residual, gamma, EPS, comm_buffer, buffer_flags, False) + + +def check_certified_cells() -> None: + """Every cell: equal to finalize + hop bit for bit, within the band of the reference, and the + same on every rank.""" + gamma = _gamma(7) + for i, cell in enumerate(entry.mnnvl_moe_finalize_hop.CELLS): + spec = cell.spec + rows = spec["rows"] + assert spec["hidden"] == HIDDEN, cell.why + seed = 100 + 13 * i + per_rank = [_moe_inputs(rows, seed, r, spec["routing"]) for r in range(WORLD)] + residual = _residual(rows, seed) + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, rows, HIDDEN) + gemm2_output, weights, index, addend = per_rank[RANK] + with validating(entry.mnnvl_moe_finalize_hop): + norm, residual_out = entry.mnnvl_moe_finalize_hop( + gemm2_output, + weights, + index, + addend, + residual, + gamma, + EPS, + comm_buffer, + buffer_flags, + ) + ref_norm, ref_residual = entry.mnnvl_moe_finalize_hop.reference( + per_rank, residual, gamma, EPS + ) + _compare(norm, residual_out, ref_norm, ref_residual, cell.why) + two_step_norm, two_step_residual = _two_step(per_rank[RANK], residual, gamma) + ranks.assert_bitwise(norm, two_step_norm, f"{cell.why}: norm_out vs finalize + hop") + ranks.assert_bitwise(residual_out, two_step_residual, f"{cell.why}: residual_out") + ranks.assert_same_on_every_rank(norm, f"{cell.why}: norm_out") + ranks.assert_same_on_every_rank(residual_out, f"{cell.why}: residual_out") + dist.barrier() + + +def check_empty_slots_and_wide_rows() -> None: + """-1 slots are skipped and FC2 rows wider than the hidden size are read at their pitch.""" + gamma = _gamma(11) + for rows, pitch, empty in ((8, HIDDEN, True), (8, HIDDEN + 128, False), (3, HIDDEN + 64, True)): + seed = 2000 + rows + pitch + per_rank = [ + _moe_inputs(rows, seed, r, "top8", pitch=pitch, empty_slots=empty) for r in range(WORLD) + ] + residual = _residual(rows, seed) + norm, residual_out = _fused(per_rank[RANK], residual, gamma) + two_step = _two_step(per_rank[RANK], residual, gamma) + where = f"rows={rows} pitch={pitch} empty slots={empty}" + ranks.assert_bitwise(norm, two_step[0], f"{where}: norm_out vs finalize + hop") + ranks.assert_bitwise(residual_out, two_step[1], f"{where}: residual_out") + ref = entry.mnnvl_moe_finalize_hop.reference(per_rank, residual, gamma, EPS) + _compare(norm, residual_out, *ref, where) + dist.barrier() + + +def check_lamport_flags_rotate() -> None: + """The fused hop rotates the workspace's flags as the plain hop does.""" + gamma = _gamma(13) + for step, rows in enumerate((8, 1, 8)): + _, buffer_flags = ws.hop_workspace(MAPPING, rows, HIDDEN) + before = ranks.lamport_flags(buffer_flags) + routing = "top8" if rows > 1 else "routed_row" + _fused(_moe_inputs(rows, 3000 + step, RANK, routing), _residual(rows, step), gamma) + after = ranks.lamport_flags(buffer_flags) + bytes_to_clear = rows * HIDDEN * WORLD * torch.bfloat16.itemsize + ranks.assert_rotated(before, after, bytes_to_clear, f"step {step}") + dist.barrier() + + +def check_interleaved_with_the_hops() -> None: + """Fused hops, plain hops and the built-in all-reduce follow each other on one workspace.""" + from tensorrt_llm.functional import AllReduceFusionOp + + gamma = _gamma(17) + issued: List[Tuple[Any, ...]] = [] + for step in range(9): + rows = (8, 1, 8)[step % 3] + seed = 4000 + step + residual = _residual(rows, seed) + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, rows, HIDDEN) + if step % 3 == 2: + per_rank_rows = [ranks.normal((rows, HIDDEN), seed * 3 + r) for r in range(WORLD)] + outputs = hop.mnnvl_hop( + per_rank_rows[RANK], residual, gamma, EPS, comm_buffer, buffer_flags, False + ) + ref = hop.mnnvl_hop.reference(per_rank_rows, residual, gamma, EPS) + elif step == 4: + per_rank_rows = [ranks.normal((rows, HIDDEN), seed * 3 + r) for r in range(WORLD)] + outputs = torch.ops.trtllm.mnnvl_fusion_allreduce( + per_rank_rows[RANK], + gamma, + residual, + EPS, + comm_buffer, + buffer_flags, + True, + None, + int(AllReduceFusionOp.RESIDUAL_RMS_NORM), + ) + ref = hop.mnnvl_hop.reference(per_rank_rows, residual, gamma, EPS) + else: + routing = "top8" if rows > 1 else "routed_row" + per_rank = [_moe_inputs(rows, seed, r, routing) for r in range(WORLD)] + outputs = _fused(per_rank[RANK], residual, gamma) + ref = entry.mnnvl_moe_finalize_hop.reference(per_rank, residual, gamma, EPS) + issued.append((step, *outputs, *ref)) + torch.cuda.synchronize() + for step, norm, residual_out, ref_norm, ref_residual in issued: + _compare(norm, residual_out, ref_norm, ref_residual, f"interleaved step {step}") + + +def check_cuda_graph_replays() -> None: + """A captured verify-layer pair and a draft-step hop replay right with fresh routing.""" + gamma = _gamma(19) + static8 = [t.clone() for t in _moe_inputs(8, 5000, RANK, "top8")] + static1 = [t.clone() for t in _moe_inputs(1, 5001, RANK, "routed_row")] + residual8 = _residual(8, 5000).clone() + + def chain(): + norm8, res8 = _fused(static8, residual8, gamma) + norm1, res1 = _fused(static1, res8[:1].contiguous(), gamma) + return norm8, res8, norm1, res1 + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + chain() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs = chain() + for replay in range(5): + seed = 5100 + 10 * replay + per_rank8 = [_moe_inputs(8, seed, r, "top8") for r in range(WORLD)] + per_rank1 = [_moe_inputs(1, seed + 1, r, "routed_row") for r in range(WORLD)] + for dst, src in zip(static8, per_rank8[RANK]): + dst.copy_(src) + for dst, src in zip(static1, per_rank1[RANK]): + dst.copy_(src) + residual8.copy_(_residual(8, seed)) + torch.cuda.synchronize() + dist.barrier() + graph.replay() + torch.cuda.synchronize() + norm8, res8, norm1, res1 = outputs + ref = entry.mnnvl_moe_finalize_hop.reference(per_rank8, residual8, gamma, EPS) + _compare(norm8, res8, *ref, f"replay {replay} verify layer") + ref = entry.mnnvl_moe_finalize_hop.reference(per_rank1, res8[:1], gamma, EPS) + _compare(norm1, res1, *ref, f"replay {replay} draft step") + ranks.assert_same_on_every_rank(norm1, f"replay {replay} draft step") + dist.barrier() + del graph + torch.cuda.synchronize() + + +def check_the_guard_refuses_an_index_out_of_range() -> None: + """An index past the FC2 rows reads other memory as an expert row: is_valid refuses it.""" + gemm2_output, weights, index, addend = _moe_inputs(8, 6000, RANK, "top8") + index = index.clone() + index[3, 2] = gemm2_output.shape[0] + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, 8, HIDDEN) + residual = _residual(8, 6000) + try: + with validating(entry.mnnvl_moe_finalize_hop): + entry.mnnvl_moe_finalize_hop( + gemm2_output, + weights, + index, + addend, + residual, + _gamma(23), + EPS, + comm_buffer, + buffer_flags, + ) + except AssertionError as exc: + assert "must index gemm2_output's rows" in str(exc), exc + else: + raise AssertionError("is_valid admitted an index past the FC2 rows") + + +CHECKS = ( + check_certified_cells, + check_empty_slots_and_wide_rows, + check_lamport_flags_rotate, + check_interleaved_with_the_hops, + check_cuda_graph_replays, + check_the_guard_refuses_an_index_out_of_range, +) + + +def _run_one_rank() -> None: + global entry, hop, ws, validating, RANK, WORLD, MAPPING + RANK, WORLD, MAPPING = ranks.init_rank() + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.comm import ( # noqa: E501 + _mnnvl_workspace, + mnnvl_hop, + mnnvl_moe_finalize_hop, + ) + + entry, hop, ws = mnnvl_moe_finalize_hop, mnnvl_hop, _mnnvl_workspace + # Run as a script, not collected: the shared _validating module is reached the plain way. + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating as _validating_cm + + validating = _validating_cm + ranks.run_checks(CHECKS, RANK) + + +if __name__ == "__main__": + if ranks.WORKER_FLAG in sys.argv: + _run_one_rank() + ranks.spawn_ranks(Path(__file__).resolve(), DEADLINE_S) + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_ranks.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_ranks.py new file mode 100644 index 000000000000..634311451185 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_ranks.py @@ -0,0 +1,221 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Rank plumbing shared by this target's MNNVL op matrices. + +A matrix is its own launcher. Run plainly, it spawns one worker per device named in +CUDA_VISIBLE_DEVICES (``spawn_ranks``) and kills them all on the first failure or at its deadline, +because a broken Lamport exchange hangs rather than raising. Run with ``--rank-worker`` it is one +rank (``init_rank``). The workers rendezvous through torch.distributed's env:// method without +MPI, and the MNNVL workspaces are built over the TP process group, as in a non-MPI serving engine. + +The c1 world, 16 ranks on 4 nodes of 4 GPUs on one MNNVL fabric, is started by hand, one worker +per GPU, for example inside a 4-node allocation: + + export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -1) + export MASTER_PORT=29511 TRTLLM_ENABLE_PDL=1 + srun -N4 --ntasks-per-node=4 python _mnnvl_hop_op_matrix.py --rank-worker + +``init_rank`` then reads the rank, the world size and the local rank from Slurm. The checks that +need an oracle only a 1 MiB one-shot all-reduce provides skip the shapes above it. + +Only torch is imported at module scope: the launcher must carry no tensorrt_llm state. +""" + +from __future__ import annotations + +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Tuple + +import torch +import torch.distributed as dist + +WORKER_FLAG = "--rank-worker" + +# Lamport data per rank up to which the built-in MNNVL all-reduce runs one-shot, with the same +# reduction order as the hops: the shapes where it serves as a bitwise oracle. +BUILTIN_ONE_SHOT_BYTES = 1 << 20 + + +def spawn_ranks(script: Path, deadline_s: int) -> None: + """Start one worker per device in CUDA_VISIBLE_DEVICES and wait for all of them. + + The workers rendezvous through env:// on a free local port. Each gets its own session, so + that on the first nonzero exit or at the deadline every session is killed. + """ + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the devices this run owns, e.g. 0,1,2,3" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size >= 2, f"CUDA_VISIBLE_DEVICES names {world_size} device(s); need 2 or more" + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + ) + # The children rendezvous through env://, not through the launching step's PMIx: drop its + # namespace, or the first hard exit trips srun's errhandler before anything is reported. + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + print(f"[launcher] spawning {world_size} ranks, rendezvous on :{port}", flush=True) + children = [ + subprocess.Popen( + [sys.executable, str(script), WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + for r in range(world_size) + ] + + def _kill_all() -> None: + for p in children: + try: + os.killpg(p.pid, signal.SIGKILL) + except ProcessLookupError: + pass + for p in children: + p.wait() + + deadline = time.monotonic() + deadline_s + live = set(range(world_size)) + failure = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {deadline_s}s (wedged)" + break + for r in sorted(live): + code = children[r].poll() + if code is None: + continue + live.discard(r) + if code != 0: + failure = f"rank {r} exited {code}" + break + time.sleep(0.5) + if failure is not None: + _kill_all() + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + + +def init_rank() -> Tuple[int, int, Any]: + """Bind this worker's GPU, join the process group, and return ``(rank, world, mapping)``. + + ``mapping`` is the pure tensor-parallel mapping over the whole world, the c1 server's. + """ + os.environ["TLLM_DISABLE_MPI"] = "1" + if "RANK" not in os.environ and "SLURM_PROCID" in os.environ: + os.environ["RANK"] = os.environ["SLURM_PROCID"] + os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"] + os.environ["LOCAL_RANK"] = os.environ["SLURM_LOCALID"] + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ.get("LOCAL_RANK", rank)) + torch.cuda.set_device(local_rank) + # device_id keeps NCCL teardown off the implicit-device path; gloo serves the host + # agreements of the workspace construction. + dist.init_process_group( + "cuda:nccl,cpu:gloo", + init_method="env://", + rank=rank, + world_size=world, + device_id=torch.device("cuda", local_rank), + ) + from tensorrt_llm.mapping import Mapping + + gpus_per_node = min(world, torch.cuda.device_count()) + mapping = Mapping(world_size=world, rank=rank, gpus_per_node=gpus_per_node, tp_size=world) + return rank, world, mapping + + +def run_checks(checks, rank: int) -> None: + """Run every check in order; a failing rank exits hard, and the launcher kills its peers.""" + for check in checks: + try: + check() + except BaseException: + import traceback + + print(f"[rank {rank}] FAILED {check.__name__}", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + # A rank that leaves a collective early wedges its peers: exit at once. + os._exit(13) + dist.barrier() + torch.cuda.synchronize() + print(f"[rank {rank}] {len(checks)} checks passed", flush=True) + sys.stdout.flush() + # Process groups and MNNVL workspaces are torn down with the process. + os._exit(0) + + +def normal(shape, seed: int, scale: float = 1.0, dtype: torch.dtype = torch.bfloat16): + """Normal values times ``scale`` in ``dtype``, the same on every rank for one seed.""" + gen = torch.Generator(device="cuda").manual_seed(seed) + return (torch.randn(shape, generator=gen, device="cuda") * scale).to(dtype) + + +def eighths(shape, seed: int, bound: int, dtype: torch.dtype = torch.bfloat16): + """Multiples of 1/8 in [-bound / 8, bound / 8], exact in BF16 and in their sums up to 32.""" + gen = torch.Generator(device="cuda").manual_seed(seed) + raw = torch.randint(-bound, bound + 1, shape, generator=gen, device="cuda") + return (raw.float() / 8.0).to(dtype) + + +_INTEGER_OF_SIZE = {1: torch.uint8, 2: torch.int16, 4: torch.int32, 8: torch.int64} + + +def bits(t: torch.Tensor) -> torch.Tensor: + """The tensor's bits as integers of its element size.""" + return t.view(_INTEGER_OF_SIZE[t.element_size()]) + + +def assert_bitwise(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + """Gate of exactly zero, on the bits (NaN payloads and signed zeros included).""" + assert out.dtype == ref.dtype, f"{where}: dtype {out.dtype} != {ref.dtype}" + assert out.shape == ref.shape, f"{where}: shape {tuple(out.shape)} != {tuple(ref.shape)}" + differ = bits(out) != bits(ref) + if bool(differ.any()): + first = tuple(int(i) for i in differ.nonzero()[0]) + raise AssertionError( + f"{where}: {int(differ.sum())} of {out.numel()} elements differ, the first at {first}: " + f"{out[first].item()} != {ref[first].item()}" + ) + + +def assert_same_on_every_rank(t: torch.Tensor, where: str) -> None: + """Every rank of the world holds the same bits in ``t`` as rank 0.""" + gathered = [torch.empty_like(t) for _ in range(dist.get_world_size())] + dist.all_gather(gathered, t.contiguous()) + for r, other in enumerate(gathered): + assert_bitwise(other, gathered[0], f"{where}: rank {r} vs rank 0") + + +def lamport_flags(buffer_flags: torch.Tensor) -> list[int]: + """The nine Lamport flag words, as host integers.""" + torch.cuda.synchronize() + return [int(v) for v in buffer_flags.view(torch.int32).tolist()] + + +def assert_rotated(before: list[int], after: list[int], bytes_to_clear: int, where: str) -> None: + """One call later the current buffer is the dirty one, of one stage, with its bytes to clear.""" + expected = [ + (before[0] + 1) % 3, + before[0], + before[2], + 1, + bytes_to_clear, + 0, + 0, + 0, + 0, + ] + assert after == expected, f"{where}: flags {after} after {before}, expected {expected}" diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_rank_job.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_rank_job.py new file mode 100644 index 000000000000..7d3e5d460f08 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_rank_job.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Run one of this target's collective entries' rank jobs from pytest, and report what it did. + +The matrix of an entry, ``__op_matrix.py`` beside this file, is its own launcher: one rank +per device, its own process group, and a deadline it enforces by killing that group, which keeps a +wedged Lamport exchange from taking the calling run down with it. This only selects the devices, +starts the launcher in a fresh interpreter (it must carry no CUDA or torch.distributed state of +its own), and turns its exit code into an assertion. + +The ranks import ``tensorrt_llm`` from where this pytest process finds it, so a source checkout +and an installed wheel both work. +""" + +from __future__ import annotations + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path + +WORLD_SIZE = 4 +"""One node's ranks. The c1 world of 16 ranks on 4 nodes is started by hand (``_mnnvl_ranks``).""" + +_LAUNCHER_GRACE_S = 300 +"""Headroom over the matrix's own deadline, so its message wins the race.""" + + +def _devices() -> str: + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible: + devices = [d for d in visible.split(",") if d.strip()] + else: + import torch + + devices = [str(i) for i in range(torch.cuda.device_count())] + assert len(devices) >= WORLD_SIZE, ( + f"this matrix runs at world size {WORLD_SIZE}; only {len(devices)} device(s) are visible" + ) + return ",".join(devices[:WORLD_SIZE]) + + +def _load_launcher_constants(launcher: Path): + """Read the matrix's module-level deadline without importing it by name.""" + spec = importlib.util.spec_from_file_location(launcher.stem, launcher) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def run(entry: str) -> None: + """Run ``__op_matrix``'s launcher over ``WORLD_SIZE`` devices.""" + launcher = Path(__file__).resolve().with_name(f"_{entry}_op_matrix.py") + env = dict(os.environ, CUDA_VISIBLE_DEVICES=_devices()) + package = importlib.util.find_spec("tensorrt_llm") + assert package is not None and package.origin is not None, "tensorrt_llm is not importable" + package_parent = str(Path(package.origin).resolve().parents[1]) + env["PYTHONPATH"] = os.pathsep.join(p for p in (package_parent, env.get("PYTHONPATH", "")) if p) + timeout = _load_launcher_constants(launcher).DEADLINE_S + _LAUNCHER_GRACE_S + completed = subprocess.run( + [sys.executable, str(launcher)], + env=env, + capture_output=True, + text=True, + timeout=timeout, + ) + if completed.returncode != 0: + raise AssertionError( + f"{launcher.name} exited {completed.returncode}\n" + f"--- stdout ---\n{completed.stdout}\n" + f"--- stderr ---\n{completed.stderr}" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_allgather_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_allgather_op_matrix.py new file mode 100644 index 000000000000..563639be3dbe --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_allgather_op_matrix.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the mnnvl_allgather entry's certification matrix. + +The matrix is ``_mnnvl_allgather_op_matrix.py`` beside this file, its own 4-rank launcher on one +node; see ``_rank_job`` for why it runs as a separate job. The 16-rank c1 world is run by hand +(``_mnnvl_ranks``). +""" + +import _rank_job +import pytest + +from tensorrt_llm._utils import get_sm_version + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the glm_moe_dsa tp16 target is certified on sm_107 only" +) + + +# The job spawns 4 ranks over the visible devices; under xdist several workers would fight for +# the same GPUs. +@pytest.mark.no_xdist +def test_mnnvl_allgather_op_matrix() -> None: + _rank_job.run("mnnvl_allgather") diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_hop_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_hop_op_matrix.py new file mode 100644 index 000000000000..6a2ea85bc49a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_hop_op_matrix.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the mnnvl_hop entry's certification matrix. + +The matrix is ``_mnnvl_hop_op_matrix.py`` beside this file, its own 4-rank launcher on one node; +see ``_rank_job`` for why it runs as a separate job. The 16-rank c1 world is run by hand +(``_mnnvl_ranks``). +""" + +import _rank_job +import pytest + +from tensorrt_llm._utils import get_sm_version + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the glm_moe_dsa tp16 target is certified on sm_107 only" +) + + +# The job spawns 4 ranks over the visible devices; under xdist several workers would fight for +# the same GPUs. +@pytest.mark.no_xdist +def test_mnnvl_hop_op_matrix() -> None: + _rank_job.run("mnnvl_hop") diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_moe_finalize_hop_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_moe_finalize_hop_op_matrix.py new file mode 100644 index 000000000000..f3254d88db23 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_moe_finalize_hop_op_matrix.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the mnnvl_moe_finalize_hop entry's certification matrix. + +The matrix is ``_mnnvl_moe_finalize_hop_op_matrix.py`` beside this file, its own 4-rank launcher +on one node; see ``_rank_job`` for why it runs as a separate job. The 16-rank c1 world is run by +hand (``_mnnvl_ranks``). +""" + +import _rank_job +import pytest + +from tensorrt_llm._utils import get_sm_version + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the glm_moe_dsa tp16 target is certified on sm_107 only" +) + + +# The job spawns 4 ranks over the visible devices; under xdist several workers would fight for +# the same GPUs. +@pytest.mark.no_xdist +def test_mnnvl_moe_finalize_hop_op_matrix() -> None: + _rank_job.run("mnnvl_moe_finalize_hop") diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_glm_rank_job.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_glm_rank_job.py new file mode 100644 index 000000000000..ede11eb3e57b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_glm_rank_job.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Run a collective entry's rank job from pytest, and report what it did. + +The GLM target's copy of ``tests/unittest/_torch/modeling_v2/comm/_rank_job.py``, under its own +module name so that the two never shadow each other in one pytest process: it starts the +``__op_matrix.py`` launcher beside this file. + +The entry's own launcher already owns the hard parts -- one rank per visible +device, its own process group, and a deadline it enforces by killing that +group, which is what keeps a wedged collective from taking the calling run +down with it. A collective that breaks *hangs* rather than raising, so the +deadline is load-bearing and the ``mpi_pool_executor`` fixture has none. + +So this does not reimplement any of that. It selects the devices, starts the +launcher in a fresh interpreter, and turns its exit code into an assertion. + +Fresh interpreter is required, not tidiness: the launcher owns the deadline +that kills a wedged collective's process sessions, and it must carry no CUDA +or torch.distributed state of its own -- a pytest process that has imported +``tensorrt_llm`` does. + +Everything here is started **by file path**: the ranks reach the catalog by +absolute import, so neither half needs a package context. This tree does not +have one to give: it is tests/, not a package. +""" + +from __future__ import annotations + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path + +WORLD_SIZE = 4 +"""The single-node world size the matrices run at; their 16-rank cells need a 4-node launch.""" + +_LAUNCHER_GRACE_S = 300 +"""Headroom over the entry's own deadline, so its message wins the race.""" + + +def _devices() -> str: + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible: + devices = [d for d in visible.split(",") if d.strip()] + else: + import torch + + devices = [str(i) for i in range(torch.cuda.device_count())] + assert len(devices) >= WORLD_SIZE, ( + f"this entry is certified at world size {WORLD_SIZE}; only " + f"{len(devices)} device(s) are visible" + ) + return ",".join(devices[:WORLD_SIZE]) + + +def _load_launcher_constants(launcher: Path): + """Read the entry's module-level budget without importing it by name. + + There is no dotted name to import it by -- this directory is not a + package -- and only its constants are wanted. Executing the file at module + scope is cheap and pulls in nothing but torch; the entry keeps its + tensorrt_llm imports inside the rank body for exactly that reason. + """ + spec = importlib.util.spec_from_file_location(launcher.stem, launcher) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def run(entry: str) -> None: + """Run ``__op_matrix``'s launcher over ``WORLD_SIZE`` devices.""" + launcher = Path(__file__).resolve().with_name(f"_{entry}_op_matrix.py") + env = dict(os.environ, CUDA_VISIBLE_DEVICES=_devices()) + + # The ranks import tensorrt_llm absolutely: the package this process + # imports, installed or not. + package = importlib.util.find_spec("tensorrt_llm") + assert package is not None and package.origin is not None, "tensorrt_llm is not importable" + package_parent = str(Path(package.origin).resolve().parents[1]) + env["PYTHONPATH"] = os.pathsep.join(p for p in (package_parent, env.get("PYTHONPATH", "")) if p) + + # Read the budget off the entry rather than restating it: an entry that + # raises its own deadline would otherwise be killed by this one first, + # and the message a reader needs ("wedged") would be lost. + constants = _load_launcher_constants(launcher) + timeout = constants.DEADLINE_S + getattr(constants, "WEDGE_CAP_S", 0) + _LAUNCHER_GRACE_S + + completed = subprocess.run( + [sys.executable, str(launcher)], + env=env, + capture_output=True, + text=True, + timeout=timeout, + ) + if completed.returncode != 0: + raise AssertionError( + f"{launcher.name} exited {completed.returncode}\n" + f"--- stdout ---\n{completed.stdout}\n" + f"--- stderr ---\n{completed.stderr}" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_hist_topk_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_hist_topk_op_matrix.py new file mode 100644 index 000000000000..328e916d0087 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_hist_topk_op_matrix.py @@ -0,0 +1,359 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for the GLM hist_topk entry (the key-sharded decode indexer top-k). + +A collective cannot be exercised in one process, so this script is its own launcher: run plainly, +it spawns one rank per device named in CUDA_VISIBLE_DEVICES, under a deadline it enforces by +killing every rank's process group. A broken exchange waits (up to the kernel's 120 s peer timeout) +rather than raising, so the deadline is what keeps it from taking the calling run down:: + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _hist_topk_op_matrix.py + +That run drives the entry's 4-rank cells. Its 16-rank cells (c1) need 16 ranks over 4 nodes: start +this file with ``--rank-worker`` on every task, one per GPU, with ``MASTER_ADDR`` / ``MASTER_PORT`` +pointing at rank 0's node and a free port. ``RANK``, ``WORLD_SIZE``, ``LOCAL_RANK`` and +``LOCAL_WORLD_SIZE`` fall back to Slurm's ``SLURM_PROCID``, ``SLURM_NTASKS``, ``SLURM_LOCALID`` and +``SLURM_NTASKS_PER_NODE``:: + + srun -N4 --ntasks-per-node=4 env MASTER_ADDR= MASTER_PORT= \\ + python _hist_topk_op_matrix.py --rank-worker + +Every rank builds the same full logits from a seed, scores nothing itself, and hands the op its own +shard (the columns past its keys are NaN, which would outrank +inf if they were read). The checks: +the cells against ``reference`` (row sets, the order within a row being unspecified), the raw +output bit-identical on all ranks, back-to-back calls through the rotating Lamport buffers, a call +without an exchange leaving the workspace untouched, CUDA-graph replays with new logits and +lengths, and the guard. + +Not a pytest module: the ``check_*`` bodies are one fixed sequence inside a single job and share +the rank's workspace. The collected entry point is ``test_modeling_v2_hist_topk_op_matrix.py``, +which starts this file and turns its exit code into an assertion. +""" + +import os +import signal +import socket +import subprocess +import sys +import time +import zlib +from pathlib import Path +from typing import Any, List, Tuple + +import torch +import torch.distributed as dist + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 1200 + +# Bound inside the rank body, never in the launcher: importing the entry pulls in tensorrt_llm, +# which must not create CUDA or torch.distributed state in the launching process. +op: Any = None +canonical_rows: Any = None +shard_size: Any = None +shard_capacity: Any = None +validating: Any = None +WORKSPACE: Any = None +RANK = 0 +WORLD = 1 + +_SPECIAL_BITS = ( + 0x7F800000, # +inf + -0x00800000, # -inf (0xFF800000) + 0x00000000, # +0 + -0x80000000, # -0 (0x80000000) + 0x7FC00000, # NaN + -0x00400000, # NaN with the sign bit (0xFFC00000) + 0x7FFFFFFF, # NaN whose key is clamped below the empty word +) + + +def _full_logits(spec: dict, seed: int) -> torch.Tensor: + """The request's logits [rows, kv_len], the same on every rank.""" + gen = torch.Generator().manual_seed(seed) + rows, kv_len, pattern = spec["rows"], spec["kv_len"], spec["pattern"] + if pattern == "ties": + logits = torch.randint(0, 64, (rows, kv_len), generator=gen).float() + elif pattern == "equal": + logits = torch.ones(rows, kv_len) + elif pattern == "saturated": + # Sign, exponent and the top two mantissa bits equal: one top-11-bit bin holds them all. + logits = 1.0 + torch.rand(rows, kv_len, generator=gen) * 2.0**-10 + else: + logits = torch.randn(rows, kv_len, generator=gen) + if pattern == "special": + values = torch.tensor(_SPECIAL_BITS, dtype=torch.int64).to(torch.int32).view(torch.float32) + for value in values: + index = torch.randint(0, kv_len, (rows, 64), generator=gen) + logits.scatter_(1, index, value.expand(rows, 64).contiguous()) + return logits.cuda() + + +def _shard(full: torch.Tensor, kv_len: int, rank: int, width: int) -> torch.Tensor: + """Rank ``rank``'s shard logits [rows, width]: its keys, then NaN.""" + size = shard_size(kv_len, WORLD) + n = min(max(kv_len - rank * size, 0), size) + shard = torch.full((full.shape[0], width), float("nan"), device="cuda") + shard[:, :n] = full[:, rank * size : rank * size + n] + return shard + + +def _call( + full: torch.Tensor, spec: dict, width: int = 0 +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """One call of the op on this rank's shard: (out, logits, kv_lens).""" + kv_len = spec["kv_len"] + logits = _shard(full, kv_len, RANK, max(width, shard_size(kv_len, WORLD))) + kv_lens = torch.tensor([kv_len], dtype=torch.int32, device="cuda") + out = torch.full((spec["rows"], 2048), -7, dtype=torch.int32, device="cuda") + with validating(op): + op( + logits, + kv_lens, + out, + WORKSPACE.comm_buffer, + WORKSPACE.state, + spec["next_n"], + WORLD, + ) + torch.cuda.synchronize() + return out, logits, kv_lens + + +def _verify(full: torch.Tensor, spec: dict, out: torch.Tensor, kv_lens: torch.Tensor, where: str): + """The row sets against ``reference``; the raw rows bit-identical on every rank.""" + width = shard_size(spec["kv_len"], WORLD) + shards: List[torch.Tensor] = [_shard(full, spec["kv_len"], r, width) for r in range(WORLD)] + expected = op.reference( + shards, kv_lens, out, WORKSPACE.comm_buffer, WORKSPACE.state, spec["next_n"], WORLD + ) + got = canonical_rows(out) + if not torch.equal(got, expected): + bad = (got != expected).any(dim=1).nonzero().flatten().tolist() + raise AssertionError(f"{where}: rows {bad} hold another set than the reference") + gathered = [torch.empty_like(out) for _ in range(WORLD)] + dist.all_gather(gathered, out) + for rank, rows in enumerate(gathered): + assert torch.equal(rows, out), f"{where}: rank {rank} wrote other bits than rank {RANK}" + assert int(WORKSPACE.state[2]) == 0, f"{where}: the op recorded a peer timeout" + + +def check_certified_cells() -> None: + ran, skipped = 0, [] + for cell in op.CELLS: + spec = cell.spec + if spec["group_size"] != WORLD: + skipped.append(cell.why) + continue + full = _full_logits(spec, zlib.crc32(cell.why.encode()) % 2**31) + out, _, kv_lens = _call(full, spec) + _verify(full, spec, out, kv_lens, cell.why) + ran += 1 + assert ran > 0, f"no cell of hist_topk.CELLS runs at {WORLD} ranks" + if RANK == 0: + print(f"[rank 0] {ran} cells at {WORLD} ranks; not run here: {skipped}", flush=True) + + +def check_back_to_back_calls_rotate_the_lamport_buffers() -> None: + """Calls with and without an exchange, of every size, one after the other: each exchange + launch takes the next of the three buffers and clears what the previous one wrote.""" + lengths = (150_000, 1_000, 90_000, 2_051, 4 * 98_304, 60_000, 30_000, 250_000, 1_500, 5_000) + patterns = ("randn", "randn", "ties", "randn", "saturated", "equal", "special", "randn") + for step, kv_len in enumerate(lengths): + spec = dict(kv_len=kv_len, pattern=patterns[step % len(patterns)], rows=8, next_n=8) + full = _full_logits(spec, 1_000 + step) + out, _, kv_lens = _call(full, spec) + _verify(full, spec, out, kv_lens, f"call {step} ({kv_len} keys)") + + +def check_rows_without_an_exchange_leave_the_workspace_alone() -> None: + state = WORKSPACE.state.clone() + buffer = WORKSPACE.comm_buffer.clone() + spec = dict(kv_len=2_000, pattern="randn", rows=8, next_n=8) + full = _full_logits(spec, 2_000) + out, _, kv_lens = _call(full, spec) + _verify(full, spec, out, kv_lens, "2,000 keys") + assert torch.equal(WORKSPACE.state, state), "a launch without an exchange moved the state" + assert torch.equal(WORKSPACE.comm_buffer, buffer), "a launch without an exchange wrote" + + +def check_cuda_graph_replay_follows_new_logits_and_lengths() -> None: + """The decode step is one captured graph: replays read the logits and the KV length in place, + and every replay is an exchange launch of its own.""" + width = shard_capacity(300_000, WORLD) + first = dict(kv_len=200_000, pattern="randn", rows=8, next_n=8) + full = _full_logits(first, 3_000) + out, logits, kv_lens = _call(full, first, width) + _verify(full, first, out, kv_lens, "graph warm-up") + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(logits, kv_lens, out, WORKSPACE.comm_buffer, WORKSPACE.state, 8, WORLD) + for step, kv_len in enumerate((180_000, 300_000, 1_900, 120_001)): + spec = dict(kv_len=kv_len, pattern="ties" if step % 2 else "randn", rows=8, next_n=8) + full = _full_logits(spec, 3_001 + step) + logits.copy_(_shard(full, kv_len, RANK, width)) + kv_lens.fill_(kv_len) + out.fill_(-7) + dist.barrier() + graph.replay() + torch.cuda.synchronize() + _verify(full, spec, out, kv_lens, f"replay {step} ({kv_len} keys)") + + +def check_guard_refuses_a_batch_of_kv_lens() -> None: + spec = dict(kv_len=100_000, pattern="randn", rows=8, next_n=8) + logits = _shard(_full_logits(spec, 4_000), spec["kv_len"], RANK, shard_size(100_000, WORLD)) + kv_lens = torch.tensor([100_000, 7], dtype=torch.int32, device="cuda") + out = torch.empty(8, 2048, dtype=torch.int32, device="cuda") + try: + with validating(op): + op(logits, kv_lens, out, WORKSPACE.comm_buffer, WORKSPACE.state, 8, WORLD) + except AssertionError as error: + assert "kv_lens" in str(error), error + else: + raise AssertionError("is_valid admitted the KV lengths of two requests") + + +CHECKS = ( + check_certified_cells, + check_back_to_back_calls_rotate_the_lamport_buffers, + check_rows_without_an_exchange_leave_the_workspace_alone, + check_cuda_graph_replay_follows_new_logits_and_lengths, + check_guard_refuses_a_batch_of_kv_lens, +) + + +def _env_int(name: str, slurm_name: str, default: int) -> int: + value = os.environ.get(name) or os.environ.get(slurm_name) + return int(value) if value else default + + +def _run_one_rank() -> int: + """Body of one rank: run every check; a failing rank exits nonzero and the launcher kills + its siblings.""" + global op, canonical_rows, shard_size, shard_capacity, validating, WORKSPACE, RANK, WORLD + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = _env_int("RANK", "SLURM_PROCID", 0) + world = _env_int("WORLD_SIZE", "SLURM_NTASKS", 1) + local_rank = _env_int("LOCAL_RANK", "SLURM_LOCALID", rank) + local_world = _env_int("LOCAL_WORLD_SIZE", "SLURM_NTASKS_PER_NODE", world) + torch.cuda.set_device(local_rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", + init_method="env://", + rank=rank, + world_size=world, + device_id=torch.device("cuda", local_rank), + ) + + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa import ( # noqa: E501 + _hist_topk, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa import ( # noqa: E501 + hist_topk as entry, + ) + from tensorrt_llm.mapping import Mapping + + # A spawned worker, not a collected test: the repo's __extra_import_path__ is not in play, so + # the shared helper is reached the plain way, inside this rank process only. + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating as _validating_cm + + op = entry.hist_topk + canonical_rows = entry.canonical_rows + shard_size = _hist_topk.shard_size + shard_capacity = _hist_topk.shard_capacity + validating = _validating_cm + RANK, WORLD = rank, world + assert 2 <= WORLD <= _hist_topk.MAX_RANKS, f"a {WORLD}-rank group" + mapping = Mapping(world_size=world, rank=rank, gpus_per_node=local_world, tp_size=world) + WORKSPACE = _hist_topk.create_hist_topk_workspace(mapping, comm=dist.group.WORLD) + + for check in CHECKS: + try: + check() + dist.barrier() + except BaseException: + import traceback + + print(f"[rank {RANK}] FAILED {check.__name__}", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + # Exit hard: a rank that leaves a collective early wedges its siblings, and the + # launcher kills the rest when it sees this exit. + os._exit(13) + torch.cuda.synchronize() + print(f"[rank {RANK}] {len(CHECKS)} checks passed", flush=True) + # Exit without tearing the process group down: these are throwaway processes. + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """Start one worker process per claimed device on this node; kill them all on the first + nonzero exit or at the deadline.""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the devices this run owns, e.g. 0,1,2,3" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size >= 2, f"CUDA_VISIBLE_DEVICES names {world_size} device(s); need at least 2" + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + LOCAL_WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + ) + # Drop the launching step's PMIx/PMI/OMPI/Slurm namespace from the children: they would + # otherwise register as participants of that step, and the first hard exit would have the + # step killed before anything is reported. + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + print(f"[launcher] spawning {world_size} ranks, rendezvous on :{port}", flush=True) + children = [ + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + for r in range(world_size) + ] + + def _kill_all() -> None: + for p in children: + try: + os.killpg(p.pid, signal.SIGKILL) + except ProcessLookupError: + pass + for p in children: + p.wait() + + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" + break + for r in sorted(live): + code = children[r].poll() + if code is None: + continue + live.discard(r) + if code != 0: + failure = f"rank {r} exited {code}" + break + time.sleep(0.5) + if failure is not None: + _kill_all() + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + + +if __name__ == "__main__": + if _WORKER_FLAG in sys.argv: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_indexer_inputs.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_indexer_inputs.py new file mode 100644 index 000000000000..e1419eafb0eb --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_indexer_inputs.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Inputs the GLM indexer entries' tests share: the indexer K cache, its slot mappings and block +table, the RoPE table, and the comparison of E4M3 rows quantized with per-row scales. + +The cache layout is the engine's: an INDEX_KEY page holds 64 tokens' 128 FP8 bytes, then their 64 +fp32 scales, viewed as [pages, 64, 1, 132] uint8; a token's bytes are found through the slot +mappings, flat byte offsets computed exactly as the DSA metadata does. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +PAGE = 64 +HEAD_DIM = 128 +TOKEN_BYTES = HEAD_DIM + 4 +SENTINEL = 0xA5 + + +def cos_sin_table(rows: int, rope_dim: int = 64, theta: float = 8.0e6) -> torch.Tensor: + """[rows, rope_dim] fp32, ``[cos | sin]`` of the default RoPE at every position.""" + exponents = torch.arange(0, rope_dim, 2, dtype=torch.float64, device="cuda") / rope_dim + inv_freq = 1.0 / (theta**exponents) + angles = torch.arange(rows, dtype=torch.float64, device="cuda")[:, None] * inv_freq[None, :] + return torch.cat([angles.cos(), angles.sin()], dim=1).float() + + +def k_cache_pool(pages: int, layout: str) -> Tuple[torch.Tensor, torch.Tensor]: + """(this layer's [pages, 64, 1, 132] uint8 view, the whole allocation), filled with a + sentinel. ``interleaved`` puts two other layers' pages between this layer's pages (a strided + page dimension); ``contiguous`` is one layer's own pool.""" + if layout == "interleaved": + whole = torch.full((pages, 3, PAGE * TOKEN_BYTES), SENTINEL, dtype=torch.uint8).cuda() + view = whole[:, 1].view(pages, PAGE, 1, TOKEN_BYTES) + assert not view.is_contiguous() + return view, whole + whole = torch.full((pages, PAGE, 1, TOKEN_BYTES), SENTINEL, dtype=torch.uint8).cuda() + return whole, whole + + +def block_row(positions: torch.Tensor, pool_pages: int, gen: torch.Generator) -> torch.Tensor: + """One request's INDEX_KEY block-table row up to the page of its last position: the pages the + positions touch map to distinct random physical pages (so no two rows share a slot), the + others to page 0.""" + pages = positions.long() // PAGE + touched = torch.unique(pages) + assert touched.numel() <= pool_pages, "the pool is smaller than the pages the rows touch" + row = torch.zeros(int(pages.max()) + 1, dtype=torch.int32) + row[touched] = torch.randperm(pool_pages, generator=gen)[: touched.numel()].to(torch.int32) + return row + + +def slot_mappings(positions: torch.Tensor, row: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """(fp8 slots, scale slots) int64 flat byte offsets of each position's K row, the metadata's + ``_compute_slot_mappings`` for one request.""" + pos = positions.long().cpu() + block_ids = row.long()[pos // PAGE] + in_page = pos % PAGE + page_stride = PAGE * TOKEN_BYTES + fp8 = block_ids * page_stride + in_page * HEAD_DIM + scale = block_ids * page_stride + PAGE * HEAD_DIM + in_page * 4 + return fp8.cuda(), scale.cuda() + + +def e4m3_step(units: torch.Tensor) -> torch.Tensor: + """Spacing of E4M3 values at magnitude ``units`` (in units of the scale).""" + exponent = torch.floor(torch.log2(units.clamp_min(2.0**-6))) + return torch.pow(2.0, exponent - 3) + + +def assert_quantized_close( + q8: torch.Tensor, + scale: torch.Tensor, + q8_ref: torch.Tensor, + scale_ref: torch.Tensor, + max_mismatch: float = 0.01, +) -> None: + """E4M3 rows with per-row scales against a reference that rounds at the same points. + + The kernel and the reference compute the same fp32 expressions, but the LayerNorm reduction + order and the MUFU rsqrt can move an fp32 intermediate by an ulp, and a bf16 rounding of that + intermediate then flips with probability of about 2^-16 per value. A flipped bf16 value moves + its E4M3 code by at most one step (and the row scale, a power of two, only when the row's + maximum crosses one). So every dequantized value is within one E4M3 step of the reference, at + the coarser of the two scales, and almost all codes match exactly. + """ + assert q8.shape == q8_ref.shape and scale.shape == scale_ref.shape + deq = q8.float() * scale + deq_ref = q8_ref.float() * scale_ref + step = torch.maximum( + scale * e4m3_step(q8.float().abs()), scale_ref * e4m3_step(q8_ref.float().abs()) + ) + excess = ((deq - deq_ref).abs() - step).clamp_min(0) + assert excess.max().item() == 0, ( + f"a value is more than one E4M3 step from the reference (by {excess.max().item():.3e})" + ) + mismatch = (q8.view(torch.uint8) != q8_ref.view(torch.uint8)).float().mean().item() + assert mismatch <= max_mismatch, f"{mismatch:.2%} of the E4M3 codes differ" diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_convert_req_index_to_global.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_convert_req_index_to_global.py new file mode 100644 index 000000000000..e81383ff5d33 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_convert_req_index_to_global.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM convert_req_index_to_global entry: integer pool rows, compared exactly. + +The block table has more rows than requests (the context rows of the engine's table come first +when ``table_layout`` is ``gen_rows``) and -1 entries past each request's pages; the TopK rows hold +request-local tokens below each request's KV length, with ``invalid`` -1 slots at the end. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.convert_req_index_to_global import ( # noqa: E501 + convert_req_index_to_global as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +assert torch.cuda.is_available(), "convert_req_index_to_global requires a CUDA device" + + +def _build(spec, seed: int) -> dict: + gen = torch.Generator().manual_seed(seed) + rows, topk, tpb = spec["rows"], spec["topk"], spec["tokens_per_block"] + width, pool = spec["table_width"], spec["pool_pages"] + pages = -(-spec["kv_len"] // tpb) + table = torch.full((rows + 2, width), -1, dtype=torch.int32) + for r in range(rows + 2): + table[r, :pages] = torch.randint(0, pool, (pages,), generator=gen, dtype=torch.int32) + if spec["table_layout"] == "gen_rows": + # The engine's table holds the context requests' rows first; the op gets the generation + # requests' rows, a view starting past them. + table = table.cuda()[2:] + else: + table = table[:rows].cuda() + # A request shorter than the TopK fills its first kv_len slots; the rest are -1. + valid = topk - spec["invalid"] + tokens = torch.full((rows, topk), -1, dtype=torch.int32) + for r in range(rows): + tokens[r, :valid] = torch.randperm(spec["kv_len"], generator=gen)[:valid].to(torch.int32) + return dict( + req_id=torch.arange(rows, dtype=torch.int32, device="cuda"), + block_table=table, + token_indices=tokens.cuda(), + block_size=tpb, + num_topk_tokens=topk, + stride_factor=spec["page_index_scale"] * tpb, + layer_id=spec["layer"], + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + ref = op.reference(**args) + out = op(**args) + torch.cuda.synchronize() + assert torch.equal(out, ref) + if cell.spec["invalid"]: + assert (out[:, -cell.spec["invalid"] :] == -1).all() + + +def test_guard_refuses_fewer_topk_columns_than_the_rows_hold(): + args = _build(op.CELLS[0].spec, seed=81) + with pytest.raises(AssertionError, match="unwritten"): + with validating(op): + op(**dict(args, num_topk_tokens=1024)) + + +def test_guard_refuses_a_request_past_the_table(): + args = _build(op.CELLS[0].spec, seed=82) + req = torch.full_like(args["req_id"], args["block_table"].shape[0]) + with pytest.raises(AssertionError, match="past the block table"): + with validating(op): + op(**dict(args, req_id=req)) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_fp8_paged_mqa_logits.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_fp8_paged_mqa_logits.py new file mode 100644 index 000000000000..56757d1a6684 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_fp8_paged_mqa_logits.py @@ -0,0 +1,202 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM fp8_paged_mqa_logits entry. + +The keys live in a real INDEX_KEY pool of 64-token pages, scattered over random physical pages; the +schedule comes from the production schedule op (``paged_mqa_logits_schedule``), as at c1. + +Most cells use integer data: Q and K are small integers (exact in E4M3), the key scales powers of +two and the weights small integers times 2^-4. Every product and partial sum is then an exactly +representable fp32 value, whatever the accumulation order, so the kernel must equal the fp64 +reference bit for bit. One cell uses UE8M0-quantized random keys and real-valued weights, where the +accumulation order is visible and a band applies. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.fp8_paged_mqa_logits import ( # noqa: E501 + fp8_paged_mqa_logits as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.paged_mqa_logits_schedule import ( # noqa: E501 + paged_mqa_logits_schedule, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../..", "."] +import _indexer_inputs as inputs # noqa: E402 +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +assert torch.cuda.is_available(), "fp8_paged_mqa_logits requires a CUDA device" + +_SPLIT = 256 + + +def _sms() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + +def _keys(count: int, data: str, gen: torch.Generator): + """(E4M3 keys [count, 128], fp32 scales [count]).""" + if data == "int": + keys = torch.randint(-3, 4, (count, inputs.HEAD_DIM), generator=gen).float() + scales = torch.pow(2.0, torch.randint(-6, 2, (count,), generator=gen).float()) + return keys.to(torch.float8_e4m3fn), scales + raw = torch.randn(count, inputs.HEAD_DIM, generator=gen) + amax = raw.abs().amax(dim=1).clamp_min(1e-4) + scales = torch.pow(2.0, torch.ceil(torch.log2(amax / 448.0))) + return (raw / scales[:, None]).to(torch.float8_e4m3fn), scales + + +def _build(spec, seed: int) -> dict: + gen = torch.Generator().manual_seed(seed) + lens = list(spec["context_lens"]) + batch, next_n, heads, d = len(lens), spec["next_n"], spec["num_heads"], spec["head_dim"] + page = spec["tokens_per_block"] + max_ctx = spec["max_context_len"] + width = -(-max_ctx // page) + pages_per_request = [-(-n // page) for n in lens] + pool_pages = sum(pages_per_request) + 3 + cache, whole = inputs.k_cache_pool(pool_pages, spec["pool_layout"]) + order = torch.randperm(pool_pages, generator=gen).to(torch.int32) + table = torch.zeros(batch, width, dtype=torch.int32) + used = 0 + flat_pages = cache.reshape(pool_pages, -1) + for b, n in enumerate(lens): + count = pages_per_request[b] + table[b, :count] = order[used : used + count] + used += count + keys, scales = _keys(count * page, spec["data"], gen) + page_bytes = torch.cat( + [ + keys.view(torch.uint8).reshape(count, page * d), + scales.view(torch.uint8).reshape(count, page * 4), + ], + dim=1, + ) + flat_pages[table[b, :count].long().cuda()] = page_bytes.cuda() + if spec["data"] == "int": + q = torch.randint(-3, 4, (batch, next_n, heads, d), generator=gen).float() + weights = torch.randint(-3, 4, (batch * next_n, heads), generator=gen).float() * 2.0**-4 + else: + q = torch.randn(batch, next_n, heads, d, generator=gen) + weights = torch.randn(batch * next_n, heads, generator=gen) * 0.05 + context_lens = torch.tensor(lens, dtype=torch.int32).cuda() + schedule = torch.empty(_sms() + 1, 2, dtype=torch.int32, device="cuda") + paged_mqa_logits_schedule(context_lens, page, _sms(), schedule) + logits = torch.full( + (batch * next_n, -(-max_ctx // _SPLIT) * _SPLIT), float("nan"), device="cuda" + ) + return dict( + q=q.to(torch.float8_e4m3fn).cuda(), + kv_cache=cache, + weights=weights.cuda(), + context_lens=context_lens, + block_table=table.cuda(), + schedule_meta=schedule, + logits=logits, + max_context_len=max_ctx, + early_kv=spec["early_kv"], + ) + + +def _compare(out: torch.Tensor, ref: torch.Tensor, lens, next_n: int, exact: bool) -> None: + """Request by request over its keys; the columns past a request's keys are unspecified.""" + for b, n in enumerate(lens): + got = out[b * next_n : (b + 1) * next_n, :n] + want = ref[b * next_n : (b + 1) * next_n, :n] + if exact: + assert torch.equal(got, want), f"request {b}: logits differ from the exact reference" + else: + # fp32 accumulation in another order: products of E4M3 values are exact, so the + # difference is the rounding of the 128-term dots and of the 32-term weighted sum, + # a few ulps of the partial sums. 64 ulps of each row's largest logit leave an order of + # magnitude of margin; a wrong head, weight or key is off by the logits themselves. + assert_within_ulp(got.contiguous(), want.contiguous(), element_ulp=64.0, rms_ulp=8.0) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + ref = op.reference(**args) + out = op(**args) + torch.cuda.synchronize() + assert out.data_ptr() == args["logits"].data_ptr(), "the op writes the caller's buffer" + assert out.shape == (args["logits"].shape[0], cell.spec["max_context_len"]) + _compare(out, ref, cell.spec["context_lens"], cell.spec["next_n"], cell.spec["data"] == "int") + + +def test_early_kv_kernel_equals_the_default_kernel_bitwise(): + """The N split and the early KV loads move work, not arithmetic: same bits on real values.""" + spec = next(c.spec for c in op.CELLS if c.spec["data"] == "randn") + args = _build(spec, seed=51) + early = op(**dict(args, early_kv=True)).clone() + default = op(**dict(args, early_kv=False)) + torch.cuda.synchronize() + n = spec["context_lens"][0] + assert torch.equal(early[:, :n], default[:, :n]) + + +def test_cuda_graph_replay_reads_new_queries_and_keys(): + """The decode path is captured: a replay follows Q, the weights and the cache in place.""" + spec = op.CELLS[0].spec + args = _build(spec, seed=52) + op(**args) # compile outside capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(**args) + fresh = _build(spec, seed=53) + args["q"].copy_(fresh["q"]) + args["weights"].copy_(fresh["weights"]) + args["kv_cache"].copy_(fresh["kv_cache"]) + args["block_table"].copy_(fresh["block_table"]) + ref = op.reference(**args) + graph.replay() + torch.cuda.synchronize() + _compare(args["logits"], ref, spec["context_lens"], spec["next_n"], exact=True) + + +def test_guard_refuses_a_block_table_too_narrow_for_the_keys(): + args = _build(op.CELLS[2].spec, seed=54) + narrow = dict( + args, + context_lens=torch.tensor([5_000], dtype=torch.int32, device="cuda"), + max_context_len=5_000, + block_table=args["block_table"][:, :10].contiguous(), + ) + with pytest.raises(AssertionError, match="cannot cover"): + with validating(op): + op(**narrow) + + +def test_guard_refuses_a_schedule_of_another_sm_count(): + args = _build(op.CELLS[2].spec, seed=55) + schedule = torch.zeros(_sms(), 2, dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match="schedule_meta"): + with validating(op): + op(**dict(args, schedule_meta=schedule)) + + +def test_guard_refuses_logits_without_room_for_whole_splits(): + args = _build(op.CELLS[2].spec, seed=56) + narrow = torch.empty(args["logits"].shape[0], args["max_context_len"] - 64, device="cuda") + with pytest.raises(AssertionError, match="whole"): + with validating(op): + op(**dict(args, logits=narrow)) + + +def test_guard_refuses_strided_weights(): + args = _build(op.CELLS[2].spec, seed=57) + strided = args["weights"].t().contiguous().t() + assert not strided.is_contiguous() + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op(**dict(args, weights=strided)) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_shard_meta.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_shard_meta.py new file mode 100644 index 000000000000..ff3d0961e9ec --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_shard_meta.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM hist_shard_meta entry: integer outputs, compared exactly. + +The request's INDEX_KEY block-table row holds distinct random page ids (with slack past the +request's pages, as the engine's table rows have), and both outputs start out as a sentinel so a +missing write shows. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa._hist_topk import ( # noqa: E501 + shard_size, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.hist_shard_meta import ( # noqa: E501 + hist_shard_meta as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +assert torch.cuda.is_available(), "hist_shard_meta requires a CUDA device" + + +def _build(spec, seed: int) -> dict: + gen = torch.Generator().manual_seed(seed) + tpb = spec["tokens_per_block"] + pages = -(-spec["kv_len"] // tpb) + 7 + row = torch.randperm(4 * pages, generator=gen)[:pages].to(torch.int32) + width = -(-spec["capacity"] // tpb) + return dict( + kv_lens=torch.tensor([spec["kv_len"]], dtype=torch.int32, device="cuda"), + block_table=row.cuda(), + shard_ctx=torch.full((1,), -7, dtype=torch.int32, device="cuda"), + shard_block_table=torch.full((1, width), -7, dtype=torch.int32, device="cuda"), + group_size=spec["group_size"], + rank=spec["rank"], + tokens_per_block=tpb, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + ref_ctx, ref_table = op.reference(**args) + ctx, table = op(**args) + torch.cuda.synchronize() + assert ctx is args["shard_ctx"] and table is args["shard_block_table"] + assert torch.equal(ctx, ref_ctx), (ctx.tolist(), ref_ctx.tolist()) + assert torch.equal(table, ref_table) + + +def test_shards_of_all_ranks_tile_the_request(): + """Independent of the reference: the 16 shards' pages, in rank order, are the request's + pages, and their key counts sum to the request's length.""" + spec = op.CELLS[0].spec + args = _build(spec, seed=61) + pages, keys = [], 0 + size = shard_size(spec["kv_len"], spec["group_size"]) + for rank in range(spec["group_size"]): + ctx, table = op(**dict(args, rank=rank)) + n = int(ctx[0]) + held = min(max(spec["kv_len"] - rank * size, 0), size) + assert n == max(held, 1) + keys += held + pages += table[0, : -(-held // spec["tokens_per_block"])].tolist() + assert keys == spec["kv_len"] + total_pages = -(-spec["kv_len"] // spec["tokens_per_block"]) + assert pages == args["block_table"][:total_pages].tolist() + + +def test_guard_refuses_a_two_dimensional_table(): + args = _build(op.CELLS[0].spec, seed=62) + with pytest.raises(AssertionError, match="one-dimensional"): + with validating(op): + op(**dict(args, block_table=args["block_table"][None, :])) + + +def test_guard_refuses_a_shard_table_narrower_than_the_shard(): + args = _build(op.CELLS[0].spec, seed=63) + with pytest.raises(AssertionError, match="shard_block_table holds"): + with validating(op): + op(**dict(args, shard_block_table=args["shard_block_table"][:, :100].contiguous())) + + +def test_guard_refuses_a_row_shorter_than_the_shard(): + args = _build(op.CELLS[1].spec, seed=64) + with pytest.raises(AssertionError, match="run past"): + with validating(op): + op(**dict(args, block_table=args["block_table"][:1000].contiguous())) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_topk_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_topk_op_matrix.py new file mode 100644 index 000000000000..d3ac1acda903 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_topk_op_matrix.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the GLM hist_topk entry's certification matrix. + +The matrix itself is ``_hist_topk_op_matrix.py`` beside this file, its own 4-rank launcher (see +``_glm_rank_job`` for why it is left intact). It drives the entry's 4-rank cells; the 16-rank cells +need the 4-node launch its docstring describes. +""" + +import _glm_rank_job +import pytest +import torch + +from tensorrt_llm._utils import get_sm_version + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +assert torch.cuda.is_available(), "hist_topk requires CUDA devices" + + +# Each case spawns its own 4 rank processes over the visible devices; under xdist several +# workers would fight for the same GPUs. +@pytest.mark.no_xdist +def test_hist_topk_op_matrix() -> None: + _glm_rank_job.run("hist_topk") diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_indexer_k_tail.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_indexer_k_tail.py new file mode 100644 index 000000000000..534f616dc582 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_indexer_k_tail.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM indexer_k_tail entry. + +The K cache is real: an INDEX_KEY pool of 64-token pages ([pages, 64, 1, 132] uint8) filled with a +sentinel, a random block-table row for the request and the metadata's slot mappings. The cache is +checked in two ways: bitwise against the op's own K outputs scattered by the reference's byte +placement (so every byte the op wrote is where it belongs, and nothing else moved), and through +those outputs against the reference arithmetic. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa import ( # noqa: E501 + _indexer_reference, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.indexer_k_tail import ( # noqa: E501 + indexer_k_tail as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../..", "."] +import _indexer_inputs as inputs # noqa: E402 +from _validating import certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +assert torch.cuda.is_available(), "indexer_k_tail requires a CUDA device" + + +def _build(spec, seed: int) -> dict: + """The op's arguments for one cell. The whole pool allocation rides along as ``_pool``.""" + gen = torch.Generator().manual_seed(seed) + rows, d = spec["rows"], spec["head_dim"] + positions = torch.arange(spec["first_position"], spec["first_position"] + rows) + if spec["k_layout"] == "fused_a": + fused = torch.randn(rows, spec["fused_a_width"], generator=gen).to(torch.bfloat16).cuda() + indexer_k = fused[:, spec["k_col"] : spec["k_col"] + d] + assert not indexer_k.is_contiguous() + else: + indexer_k = torch.randn(rows, d, generator=gen).to(torch.bfloat16).cuda() + ln_weight = (1.0 + 0.1 * torch.randn(d, generator=gen)).to(torch.bfloat16).cuda() + ln_bias = (0.1 * torch.randn(d, generator=gen)).to(torch.bfloat16).cuda() + cache, whole = inputs.k_cache_pool(spec["pages"], spec["pool_layout"]) + row = inputs.block_row(positions, spec["pages"], gen) + slot_fp8, slot_scale = inputs.slot_mappings(positions, row) + if spec["negative_slot"]: + slot_fp8[2] = -1 + return dict( + _pool=whole, + indexer_k=indexer_k, + positions=positions.to(torch.int32).cuda(), + cos_sin=inputs.cos_sin_table(int(positions[-1]) + 1, spec["rope_dim"]), + ln_weight=ln_weight, + ln_bias=ln_bias, + eps=spec["eps"], + rope_dim=spec["rope_dim"], + interleave=spec["interleave"], + use_ue8m0=spec["use_ue8m0"], + k_fp8=torch.empty(rows, d, dtype=torch.float8_e4m3fn, device="cuda"), + k_scale=torch.empty(rows, 1, dtype=torch.float32, device="cuda"), + k_cache=cache, + slot_mapping_fp8=slot_fp8, + slot_mapping_scale=slot_scale, + num_cache_tokens=spec["cache_rows"], + use_pdl=spec["use_pdl"], + launch_dependents=spec["launch_dependents"], + prewait=spec["prewait"], + ) + + +def _check(args: dict, out, ref, cache_before: torch.Tensor, whole_before=None, whole=None): + k8, k_scale, cache = out + k8_ref, k_scale_ref, _ = ref + # The arithmetic: within one E4M3 step of the reference, almost every code exact. + inputs.assert_quantized_close(k8, k_scale, k8_ref, k_scale_ref) + # The placement: the op's own rows, scattered by the reference rule into the cache as it was, + # are the cache byte for byte. + expected = cache_before.clone() + _indexer_reference.write_cache_rows( + expected, + k8, + k_scale, + args["slot_mapping_fp8"], + args["slot_mapping_scale"], + args["num_cache_tokens"], + ) + assert torch.equal(cache.view(torch.uint8), expected.view(torch.uint8)), "cache bytes differ" + if whole is not None: + # The other layers' pages of an interleaved pool are untouched. + assert torch.equal(whole[:, 0], whole_before[:, 0]) + assert torch.equal(whole[:, 2], whole_before[:, 2]) + + +def _op_args(args: dict) -> dict: + return {k: v for k, v in args.items() if not k.startswith("_")} + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + whole = args.pop("_pool") if cell.spec["pool_layout"] == "interleaved" else None + args = _op_args(args) + cache_before = args["k_cache"].clone() + whole_before = whole.clone() if whole is not None else None + ref = op.reference(**args) + out = op(**args) + torch.cuda.synchronize() + assert out[0] is args["k_fp8"] and out[1] is args["k_scale"] and out[2] is args["k_cache"] + _check(args, out, ref, cache_before, whole_before, whole) + + +def test_rows_past_num_cache_tokens_stay_out_of_the_cache(): + """Rows at and past ``num_cache_tokens`` get K outputs and no cache bytes.""" + spec = dict(op.CELLS[0].spec, cache_rows=2) + args = _op_args(_build(spec, seed=11)) + before = args["k_cache"].clone() + op(**args) + torch.cuda.synchronize() + changed = (args["k_cache"].view(torch.uint8) != before.view(torch.uint8)).sum().item() + # Two rows: 128 FP8 bytes and 4 scale bytes each, at most (a value may equal the sentinel). + assert 0 < changed <= 2 * (128 + 4) + + +def test_cuda_graph_replay_follows_new_inputs(): + """The decode path is captured: a replay reads the K rows, positions and slots in place.""" + spec = op.CELLS[0].spec + args = _op_args(_build(spec, seed=21)) + op(**args) # compile outside capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(**args) + # Positions inside the captured cos/sin table, rows of another request map. + fresh = _build(dict(spec, first_position=spec["first_position"] - 8), seed=22) + args["indexer_k"].copy_(fresh["indexer_k"]) + args["positions"].copy_(fresh["positions"]) + args["slot_mapping_fp8"].copy_(fresh["slot_mapping_fp8"]) + args["slot_mapping_scale"].copy_(fresh["slot_mapping_scale"]) + cache_before = args["k_cache"].clone() + ref = op.reference(**args) + graph.replay() + torch.cuda.synchronize() + _check(args, (args["k_fp8"], args["k_scale"], args["k_cache"]), ref, cache_before) + + +def test_guard_refuses_a_wider_cos_sin_table(): + args = _op_args(_build(op.CELLS[0].spec, seed=32)) + wide = torch.cat([args["cos_sin"], args["cos_sin"]], dim=1) + with pytest.raises(AssertionError, match="cos_sin rows"): + with validating(op): + op(**dict(args, cos_sin=wide)) + + +def test_guard_refuses_a_position_past_the_table(): + args = _op_args(_build(op.CELLS[0].spec, seed=33)) + positions = args["positions"].clone() + positions[-1] = args["cos_sin"].shape[0] + with pytest.raises(AssertionError, match="past the cos/sin table"): + with validating(op): + op(**dict(args, positions=positions)) + + +def test_guard_refuses_a_strided_k_output(): + args = _op_args(_build(op.CELLS[0].spec, seed=34)) + wide = torch.empty(args["k_fp8"].shape[0], 256, dtype=torch.float8_e4m3fn, device="cuda") + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op(**dict(args, k_fp8=wide[:, :128])) + + +def test_guard_refuses_a_slot_past_the_cache(): + args = _op_args(_build(op.CELLS[0].spec, seed=35)) + slots = args["slot_mapping_fp8"].clone() + slots[0] = args["k_cache"].numel() - 64 + with pytest.raises(AssertionError, match="writes past the cache"): + with validating(op): + op(**dict(args, slot_mapping_fp8=slots)) + + +def test_guard_refuses_more_cache_rows_than_k_rows(): + args = _op_args(_build(op.CELLS[0].spec, seed=36)) + with pytest.raises(AssertionError, match="num_cache_tokens"): + with validating(op): + op(**dict(args, num_cache_tokens=args["indexer_k"].shape[0] + 1)) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_indexer_prologue.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_indexer_prologue.py new file mode 100644 index 000000000000..e702c25a0f32 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_indexer_prologue.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM indexer_prologue entry. + +Q is the wq_b output ([N, 32 * 128] bf16); K and the weights are column views of a fused-A output +row block, as the target hands them over. The cache is the engine's INDEX_KEY layout (see +``_indexer_inputs``); its bytes are checked against the op's own K rows scattered by the reference +rule, the outputs against the reference arithmetic. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa import ( # noqa: E501 + _indexer_reference, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.indexer_prologue import ( # noqa: E501 + indexer_prologue as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../..", "."] +import _indexer_inputs as inputs # noqa: E402 +from _validating import certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +assert torch.cuda.is_available(), "indexer_prologue requires a CUDA device" + + +def _build(spec, seed: int) -> dict: + gen = torch.Generator().manual_seed(seed) + rows, heads, d = spec["rows"], spec["n_heads"], spec["head_dim"] + positions = torch.arange(spec["first_position"], spec["first_position"] + rows) + q = torch.randn(rows, heads * d, generator=gen).to(torch.bfloat16).cuda() + if spec["k_layout"] == "fused_a": + fused = torch.randn(rows, spec["fused_a_width"], generator=gen).to(torch.bfloat16).cuda() + indexer_k = fused[:, spec["k_col"] : spec["k_col"] + d] + weights = fused[:, spec["k_col"] + d : spec["k_col"] + d + heads] + else: + indexer_k = torch.randn(rows, d, generator=gen).to(torch.bfloat16).cuda() + weights = torch.randn(rows, heads, generator=gen).to(torch.bfloat16).cuda() + cache = slot_fp8 = slot_scale = None + if spec["write_cache"]: + cache, _ = inputs.k_cache_pool(spec["pages"], spec["pool_layout"]) + row = inputs.block_row(positions, spec["pages"], gen) + slot_fp8, slot_scale = inputs.slot_mappings(positions, row) + return dict( + indexer_k=indexer_k, + weights=weights, + q=q.view(rows, heads, d), + positions=positions.to(torch.int32).cuda(), + cos_sin=inputs.cos_sin_table(int(positions[-1]) + 1, spec["rope_dim"]), + ln_weight=(1.0 + 0.1 * torch.randn(d, generator=gen)).to(torch.bfloat16).cuda(), + ln_bias=(0.1 * torch.randn(d, generator=gen)).to(torch.bfloat16).cuda(), + eps=spec["eps"], + weight_scale=spec["weight_scale"], + rope_dim=spec["rope_dim"], + interleave=spec["interleave"], + use_ue8m0=spec["use_ue8m0"], + k_cache=cache, + slot_mapping_fp8=slot_fp8, + slot_mapping_scale=slot_scale, + num_cache_tokens=spec["cache_rows"] if spec["write_cache"] else 0, + use_pdl=spec["use_pdl"], + launch_dependents=spec["launch_dependents"], + prewait=spec["prewait"], + ) + + +def _check(args, out, ref, cache_before): + q8, k8, k_scale, w, q_scale = out + q8_ref, k8_ref, k_scale_ref, w_ref, q_scale_ref, _ = ref + rows, heads, d = args["q"].shape + assert q8.shape == (rows, heads, d) and q_scale.shape == (rows, heads, 1) + assert w.shape == (rows, heads) and w.dtype == torch.float32 + inputs.assert_quantized_close(q8, q_scale, q8_ref, q_scale_ref) + inputs.assert_quantized_close(k8, k_scale, k8_ref, k_scale_ref) + # weights = (w * q_scale) * weight_scale, two fp32 roundings: dividing the Q scale back out + # leaves the input weight times the constant within two ulps whichever scale either side + # picked. + torch.testing.assert_close( + w / q_scale[..., 0], + w_ref / q_scale_ref[..., 0], + rtol=2 * 2.0**-23, + atol=0.0, + ) + if args["k_cache"] is not None: + expected = cache_before.clone() + _indexer_reference.write_cache_rows( + expected, + k8, + k_scale, + args["slot_mapping_fp8"], + args["slot_mapping_scale"], + args["num_cache_tokens"], + ) + assert torch.equal(args["k_cache"], expected), "cache bytes differ" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + cache_before = args["k_cache"].clone() if args["k_cache"] is not None else None + ref = op.reference(**args) + out = op(**args) + torch.cuda.synchronize() + _check(args, out, ref, cache_before) + + +def test_q_part_matches_the_k_tail_launch_of_the_same_kernel(): + """The K part does not depend on whether the Q part is compiled in: the whole prologue and the + K tail write the same K bytes for the same inputs.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.indexer_k_tail import ( # noqa: E501 + indexer_k_tail, + ) + + args = _build(op.CELLS[0].spec, seed=41) + _, k8, k_scale, _, _ = op(**dict(args, k_cache=None, num_cache_tokens=0)) + tail_k8 = torch.empty_like(k8) + tail_scale = torch.empty_like(k_scale) + indexer_k_tail( + indexer_k=args["indexer_k"], + positions=args["positions"], + cos_sin=args["cos_sin"], + ln_weight=args["ln_weight"], + ln_bias=args["ln_bias"], + eps=args["eps"], + rope_dim=args["rope_dim"], + interleave=args["interleave"], + use_ue8m0=args["use_ue8m0"], + k_fp8=tail_k8, + k_scale=tail_scale, + k_cache=args["k_cache"], + slot_mapping_fp8=args["slot_mapping_fp8"], + slot_mapping_scale=args["slot_mapping_scale"], + num_cache_tokens=args["num_cache_tokens"], + use_pdl=args["use_pdl"], + launch_dependents=args["launch_dependents"], + prewait=args["prewait"], + ) + torch.cuda.synchronize() + assert torch.equal(k8.view(torch.uint8), tail_k8.view(torch.uint8)) + assert torch.equal(k_scale, tail_scale) + + +def test_guard_refuses_weights_of_another_head_count(): + args = _build(op.CELLS[0].spec, seed=42) + with pytest.raises(AssertionError, match="weights must be"): + with validating(op): + op(**dict(args, weights=args["weights"][:, :16])) + + +def test_guard_refuses_a_cache_write_without_slots(): + args = _build(op.CELLS[0].spec, seed=43) + with pytest.raises(AssertionError, match="both slot mappings"): + with validating(op): + op(**dict(args, slot_mapping_scale=None)) + + +def test_guard_refuses_a_non_unit_inner_stride_q(): + args = _build(op.CELLS[0].spec, seed=44) + rows, heads, d = args["q"].shape + strided = torch.empty(rows, heads, 2 * d, dtype=torch.bfloat16, device="cuda")[..., ::2] + with pytest.raises(AssertionError, match="unit inner stride"): + with validating(op): + op(**dict(args, q=strided)) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_paged_mqa_logits_schedule.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_paged_mqa_logits_schedule.py new file mode 100644 index 000000000000..cbbc5b6b61ba --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_paged_mqa_logits_schedule.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM paged_mqa_logits_schedule entry: the schedule is integers, compared exactly. + +Every cell runs at the device's SM count, the count the logits kernel launches its CTAs with. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.paged_mqa_logits_schedule import ( # noqa: E501 + paged_mqa_logits_schedule as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +assert torch.cuda.is_available(), "paged_mqa_logits_schedule requires a CUDA device" + + +def _sms() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + +def _build(spec) -> dict: + sms = _sms() + return dict( + context_lens=torch.tensor(spec["context_lens"], dtype=torch.int32, device="cuda"), + block_kv=spec["block_kv"], + num_sms=sms, + out=torch.full((sms + 1, 2), -7, dtype=torch.int32, device="cuda"), + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args = _build(cell.spec) + ref = op.reference(**args) + out = op(**args) + torch.cuda.synchronize() + assert out is args["out"] + assert torch.equal(out, ref), f"schedule differs:\n{out.tolist()}\nexpected\n{ref.tolist()}" + + +def test_the_partition_covers_every_split_once(): + """Independent of the reference: consecutive starts tile the requests' splits.""" + lens = [3_000, 70_000, 129, 40_000] + args = _build(dict(context_lens=lens, block_kv=64)) + out = op(**args).tolist() + splits = [-(-n // 256) for n in lens] + flat = [(r, s) for r, count in enumerate(splits) for s in range(count)] + position = {unit: i for i, unit in enumerate(flat)} + starts = [position.get(tuple(row), len(flat)) for row in out] + assert starts[0] == 0 and starts[-1] == len(flat) + assert all(a <= b for a, b in zip(starts, starts[1:])) + + +def test_guard_refuses_another_sm_count(): + args = _build(op.CELLS[0].spec) + sms = args["num_sms"] - 1 + with pytest.raises(AssertionError, match="SM count"): + with validating(op): + op(**dict(args, num_sms=sms, out=torch.empty(sms + 1, 2, dtype=torch.int32).cuda())) + + +def test_guard_refuses_two_dimensional_lengths(): + args = _build(op.CELLS[4].spec) + with pytest.raises(AssertionError, match="num_requests"): + with validating(op): + op(**dict(args, context_lens=args["context_lens"].view(2, 2))) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_stash_last_accepted_rows.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_stash_last_accepted_rows.py new file mode 100644 index 000000000000..c8593238d1bd --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_stash_last_accepted_rows.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM stash_last_accepted_rows entry: a row gather, compared exactly. + +The destination is either its own [requests, 2048] buffer or rows of a wider shared buffer whose +other rows must stay as they were. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.stash_last_accepted_rows import ( # noqa: E501 + stash_last_accepted_rows as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +assert torch.cuda.is_available(), "stash_last_accepted_rows requires a CUDA device" + + +def _build(spec, seed: int): + gen = torch.Generator().manual_seed(seed) + requests, next_n, n_cols = spec["requests"], spec["next_n"], spec["n_cols"] + gen_topk = torch.randint(-1, 1 << 20, (requests * next_n, n_cols), generator=gen) + shared = torch.full((requests + 3, n_cols), -7, dtype=torch.int32, device="cuda") + dst = shared[:requests] if spec["dst"] == "wide" else torch.full_like(shared[:requests], -7) + args = dict( + gen_topk=gen_topk.to(torch.int32).cuda(), + num_accepted=torch.tensor(spec["accepted"], dtype=torch.int32, device="cuda"), + dst=dst, + next_n=next_n, + ) + return args, shared + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args, shared = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + ref = op.reference(**args) + out = op(**args) + torch.cuda.synchronize() + assert out is args["dst"] + assert torch.equal(out, ref) + if cell.spec["dst"] == "wide": + assert (shared[cell.spec["requests"] :] == -7).all(), "rows past the requests moved" + + +def test_rows_follow_the_accepted_counts_in_place_under_replay(): + """Captured once, replayed with new counts: the counts are read on the device.""" + args, _ = _build(op.CELLS[0].spec, seed=71) + op(**args) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(**args) + for accepted in (1, 6, 8): + args["num_accepted"].fill_(accepted) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(args["dst"][0], args["gen_topk"][accepted - 1]) + + +def test_guard_refuses_a_topk_of_another_row_count(): + args, _ = _build(op.CELLS[0].spec, seed=72) + with pytest.raises(AssertionError, match="gen_topk must be"): + with validating(op): + op(**dict(args, gen_topk=args["gen_topk"][:4])) + + +def test_guard_refuses_a_strided_destination(): + args, _ = _build(op.CELLS[0].spec, seed=73) + wide = torch.zeros(1, 2 * 2048, dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match="unit column strides"): + with validating(op): + op(**dict(args, dst=wide[:, ::2])) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/embedding/test_modeling_v2_fused_embedding.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/embedding/test_modeling_v2_fused_embedding.py new file mode 100644 index 000000000000..555e0509d8dc --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/embedding/test_modeling_v2_fused_embedding.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM-5.2 tp16 fused_embedding entry (sm_107). + +A gather, so everything is gated bitwise: the cells against the entry's +reference, and the op against the Embedding module's own masked F.embedding +chain, which it replaces. Then CUDA-graph replay, id shapes, and the inputs the +op refuses. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.embedding.fused_embedding import ( # noqa: E501 + HIDDEN, + VOCAB, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.embedding.fused_embedding import ( # noqa: E501 + fused_embedding as op, +) +from tensorrt_llm._torch.modules.embedding import pre_comm_embedding_ops +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +assert torch.cuda.is_available(), "fused_embedding requires a CUDA device" + +#: model.embed_tokens.weight's standard deviation in the checkpoint. +TABLE_STD = 0.0076 + + +def _seed(text: str) -> int: + return zlib.crc32(text.encode()) % 2**31 + + +@pytest.fixture(scope="module") +def table(): + """The replicated [154880, 6144] bf16 table (1.9 GB), built once per module.""" + g = torch.Generator(device="cuda").manual_seed(_seed("table")) + weight = torch.randn(VOCAB, HIDDEN, generator=g, device="cuda", dtype=torch.bfloat16) + yield weight.mul_(TABLE_STD) + del weight + torch.cuda.empty_cache() + + +def _ids(spec, seed): + """`random`: uniform in-vocab ids. `edges`: the first and last rows, the -1 + padding, ids at and past the vocab, the id type's extremes, then random.""" + g = torch.Generator(device="cuda").manual_seed(seed) + ids = torch.randint(0, VOCAB, (spec["tokens"],), generator=g, device="cuda") + if spec["ids"] == "edges": + if spec["ids_dtype"] == torch.int64: + extremes = [2**33, -(2**40)] + else: + extremes = [2**31 - 1, -(2**31)] + edges = torch.tensor([-1, 0, VOCAB - 1, VOCAB, VOCAB + 1000, *extremes], device="cuda") + ids[: edges.numel()] = edges + return ids.to(spec["ids_dtype"]) + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Bitwise, compared as raw 16-bit patterns: rows are copies, so there is + no tolerance to give, and a raw compare also holds for -0.0 and NaN rows.""" + assert out.dtype == ref.dtype and out.shape == ref.shape, (out.shape, ref.shape) + assert torch.equal(out.view(torch.int16), ref.view(torch.int16)), "rows differ" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying, table) -> None: # noqa: F811 — fixture by name + certifying(op) + ids = _ids(cell.spec, seed=_seed(cell.why)) + before = ids.clone() + out = op(ids, table) + assert out.shape == (*ids.shape, HIDDEN) and out.dtype == table.dtype + assert out.is_contiguous() + assert torch.equal(ids, before), "the ids are read-only" + _compare(out, op.reference(ids, table)) + + +@pytest.mark.parametrize("ids_dtype", [torch.int32, torch.int64], ids=["int32", "int64"]) +def test_bitwise_equal_to_the_embedding_module_chain(ids_dtype, table) -> None: + """The masked F.embedding chain the Embedding module runs for an unsharded + table, with its out-of-range guard: the op computes exactly that.""" + ids = _ids(dict(tokens=64, ids_dtype=ids_dtype, ids="edges"), seed=_seed("chain")) + chain = pre_comm_embedding_ops( + ids, + table, + tp_size=1, + tp_rank=0, + tp_mode=None, + vocab_start_index=0, + vocab_end_index=VOCAB, + gather_output=False, + padding_size=0, + ) + _compare(op(ids, table), chain) + + +def test_cuda_graph_capture_and_replay(table) -> None: + """The verify forward's lookup runs inside the step's CUDA graph: replay + reads the ids the buffer holds at replay time.""" + spec = op.CELLS[0].spec + ids = _ids(spec, seed=_seed("graph")) + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(2): + op(ids, table) + torch.cuda.current_stream().wait_stream(side) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(ids, table) + + for seed, kind in (("graph random", "random"), ("graph edges", "edges")): + ids.copy_(_ids(dict(spec, ids=kind), seed=_seed(seed))) + graph.replay() + torch.cuda.synchronize() + _compare(captured, op.reference(ids, table)) + + +def test_the_output_follows_the_id_shape(table) -> None: + """Any id shape gives `ids.shape + (hidden,)`; strided ids are made + contiguous first; no ids give an empty result without a launch.""" + ids = _ids(dict(tokens=12, ids_dtype=torch.int64, ids="edges"), seed=_seed("shape")) + flat = op(ids, table) + grid = op(ids.view(3, 4), table) + assert grid.shape == (3, 4, HIDDEN) + _compare(grid.view(12, HIDDEN), flat) + _compare(op(ids[::2], table), flat[::2].contiguous()) + assert op(ids[:0], table).shape == (0, HIDDEN) + + +def _small_table(dtype=torch.bfloat16, hidden=64): + return torch.randn(32, hidden, device="cuda").to(dtype) + + +def _small_ids(): + return torch.arange(4, device="cuda", dtype=torch.int32) + + +@pytest.mark.parametrize( + "make", + [ + lambda: (_small_ids().float(), _small_table()), + lambda: (_small_ids().cpu(), _small_table()), + lambda: (_small_ids(), _small_table(torch.float16)), + lambda: (_small_ids(), _small_table(hidden=128)[:, :64]), + lambda: (_small_ids(), _small_table(hidden=60)), + ], + ids=["float ids", "cpu ids", "fp16 table", "strided table", "hidden 60"], +) +def test_inputs_outside_the_domain_raise(make) -> None: + """Loud, never a silently wrong answer: these are why `is_valid` has + nothing left to refuse.""" + input_ids, weight = make() + with pytest.raises(RuntimeError): + with validating(op): + op(input_ids, weight) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_dense_mlp_small_m.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_dense_mlp_small_m.py new file mode 100644 index 000000000000..98184625ba61 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_dense_mlp_small_m.py @@ -0,0 +1,246 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM-5.2 tp16 dense_mlp_small_m entry (sm_107). + +The cells gate the kernels against the entry's reference within a derived band. +The rest is what a band cannot show: a row's result does not depend on the +other rows, repeats are bitwise equal, a call chained right behind another one +reads its input only after the first call wrote it (programmatic dependent +launch), CUDA-graph replay, and the shape domain the op serves and refuses. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.dense_mlp_small_m import ( # noqa: E501 + HIDDEN, + INTER_PER_RANK, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.dense_mlp_small_m import ( # noqa: E501 + dense_mlp_small_m as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +assert torch.cuda.is_available(), "dense_mlp_small_m requires a CUDA device" + +#: The c1 cell: the shape and scales every non-cell test below reuses. +C1 = op.CELLS[0] + + +def _seed(text: str) -> int: + return zlib.crc32(text.encode()) % 2**31 + + +def _build(spec, seed): + """x [rows, 6144], gate_up_weight [2 * 768, 6144] and down_weight [6144, 768], bf16. + + Normal draws at the cell's scales; x is a unit-RMS row times the norm's + gamma, as post_attention_layernorm hands it over. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + x = torch.randn(spec["rows"], HIDDEN, generator=g, device="cuda") * spec["x_scale"] + gate_up = torch.randn(2 * INTER_PER_RANK, HIDDEN, generator=g, device="cuda") + down = torch.randn(HIDDEN, INTER_PER_RANK, generator=g, device="cuda") + return ( + x.to(torch.bfloat16), + (gate_up * spec["gate_up_std"]).to(torch.bfloat16), + (down * spec["down_std"]).to(torch.bfloat16), + ) + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Summation order and the fast sigmoid, nothing else. + + Kernel and reference take the same bf16 operands and round at the same + three points (gate and up sums, h, out). They differ in the fp32 order of + the 6144- and 768-long sums (about 1e-6 relative) and in the kernel's fast + exp and division inside the sigmoid (a few fp32 ulp; tens at the saturated + cell's |gate| of 40). Either can carry a value across a bf16 rounding + boundary. An output that flips moves by one bf16 ulp of itself: at most + 2^-7 of its row's largest magnitude, 2 units of the 2^-8 band unit. An h + that flips moves each output by 2^-7 |h_k w_nk|, a few hundredths of a + unit. Flips are rare (the gaps above are 1e-3 of a bf16 ulp), so the + relative RMS sits well under one unit. + + So: 4 units per element against the row scale, 2 units of relative RMS -- + twice the single-flip bound, and three orders below a wrong product (a + swapped gate/up, a lost row, a missing SwiGLU), which misses by O(1). + """ + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + x, gate_up_weight, down_weight = _build(cell.spec, seed=_seed(cell.why)) + before = x.clone() + out = op(x, gate_up_weight, down_weight) + assert out.shape == (x.shape[0], HIDDEN) and out.dtype == torch.bfloat16 + assert out.is_contiguous() + assert torch.equal(x, before), "the op must not modify x" + _compare(out, op.reference(x, gate_up_weight, down_weight)) + + +def test_a_row_does_not_see_the_other_rows() -> None: + """Row m's result is bitwise the same with 1 to 8 rows present. + + The empty token columns of a short call read as zero and are never + stored, and the arithmetic of a row does not depend on the other rows. + """ + x, gate_up_weight, down_weight = _build(C1.spec, seed=_seed("rows")) + full = op(x, gate_up_weight, down_weight) + for rows in range(1, 8): + part = op(x[:rows].contiguous(), gate_up_weight, down_weight) + assert torch.equal(part, full[:rows]), f"{rows} rows differ from the 8-row call" + + +def test_repeats_are_bitwise_equal() -> None: + """Fixed summation orders inside a CTA, across warps and across the cluster.""" + x, gate_up_weight, down_weight = _build(C1.spec, seed=_seed("repeat")) + first = op(x, gate_up_weight, down_weight) + for _ in range(4): + assert torch.equal(op(x, gate_up_weight, down_weight), first) + + +def _chain(x, layer_a, layer_b): + out_a = op(x, *layer_a) + return out_a, op(out_a, *layer_b) + + +def test_a_chained_call_reads_its_input_after_the_wait() -> None: + """The second of two back-to-back calls reads what the first one wrote. + + With PDL the down kernel of the first call releases the gate/up kernel of + the second before it computes its output; that kernel stages its weights + early and must read x only after its grid-dependency wait. Every iteration + draws a new x, so a read of the previous iteration's output (often at the + same address) or of an unwritten buffer shows up against a standalone + recompute of the second call after a device sync. Eager and graph-replayed. + """ + x, *layer_a = _build(C1.spec, seed=_seed("chain a")) + _, *layer_b = _build(C1.spec, seed=_seed("chain b")) + g = torch.Generator(device="cuda").manual_seed(_seed("chain x")) + + for _ in range(8): + x.copy_(torch.randn(x.shape, generator=g, device="cuda") * C1.spec["x_scale"]) + out_a, out_b = _chain(x, layer_a, layer_b) + torch.cuda.synchronize() + assert torch.equal(out_b, op(out_a.clone(), *layer_b)), "eager chain read a stale input" + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + _chain(x, layer_a, layer_b) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out_a, out_b = _chain(x, layer_a, layer_b) + for _ in range(8): + x.copy_(torch.randn(x.shape, generator=g, device="cuda") * C1.spec["x_scale"]) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out_b, op(out_a.clone(), *layer_b)), "replayed chain read a stale input" + + +def test_cuda_graph_capture_and_replay() -> None: + """The c1 forward runs inside one CUDA graph: replay with new x contents + gives the eager result bitwise, and the reference within the band.""" + x, gate_up_weight, down_weight = _build(C1.spec, seed=_seed("graph")) + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(2): + op(x, gate_up_weight, down_weight) + torch.cuda.current_stream().wait_stream(side) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(x, gate_up_weight, down_weight) + + fresh, _, _ = _build(C1.spec, seed=_seed("graph replay")) + x.copy_(fresh) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(captured, op(x, gate_up_weight, down_weight)) + _compare(captured, op.reference(x, gate_up_weight, down_weight)) + + +@pytest.mark.parametrize( + "rows, hidden, inter, served", + [ + (1, HIDDEN, INTER_PER_RANK, True), + (8, HIDDEN, INTER_PER_RANK, True), + (0, HIDDEN, INTER_PER_RANK, False), + (9, HIDDEN, INTER_PER_RANK, False), + (8, HIDDEN, 2 * INTER_PER_RANK, False), + (8, HIDDEN, 4 * INTER_PER_RANK, False), + (8, HIDDEN, 16 * INTER_PER_RANK, False), + (8, 7168, INTER_PER_RANK, False), + ], + ids=["1 row", "8 rows", "0 rows", "9 rows", "tp8", "tp4", "tp1", "hidden 7168"], +) +def test_the_supported_predicate(rows, hidden, inter, served) -> None: + """The host-side answer the caller routes on: 1-8 rows of the TP16 slice + only. TP4 is what a dense MLP capped at the node's 4 GPUs would hand over.""" + assert torch.ops.trtllm.glm_dense_mlp_small_m_supported(rows, hidden, inter) is served + + +def _strided_x(): + x, w13, w2 = _build(C1.spec, seed=_seed("strided")) + wide = torch.zeros(x.shape[0], 2 * HIDDEN, dtype=x.dtype, device=x.device) + wide[:, :HIDDEN] = x + return wide[:, :HIDDEN], w13, w2 + + +def _misaligned_x(): + x, w13, w2 = _build(C1.spec, seed=_seed("misaligned")) + flat = torch.zeros(x.numel() + 1, dtype=x.dtype, device=x.device) + shifted = flat[1:].view(x.shape) + shifted.copy_(x) + return shifted, w13, w2 + + +def _too_many_rows(): + x, w13, w2 = _build(dict(C1.spec, rows=9), seed=_seed("9 rows")) + return x, w13, w2 + + +def _tp8_slice(): + x, _, _ = _build(C1.spec, seed=_seed("tp8")) + w13 = torch.zeros(4 * INTER_PER_RANK, HIDDEN, dtype=torch.bfloat16, device="cuda") + w2 = torch.zeros(HIDDEN, 2 * INTER_PER_RANK, dtype=torch.bfloat16, device="cuda") + return x, w13, w2 + + +def _fp16_x(): + x, w13, w2 = _build(C1.spec, seed=_seed("fp16")) + return x.to(torch.float16), w13, w2 + + +def _mismatched_down(): + x, w13, w2 = _build(C1.spec, seed=_seed("mismatch")) + return x, w13, w2[:, : INTER_PER_RANK // 2].contiguous() + + +@pytest.mark.parametrize( + "make", + [_strided_x, _misaligned_x, _too_many_rows, _tp8_slice, _fp16_x, _mismatched_down], + ids=lambda f: f.__name__.lstrip("_"), +) +def test_inputs_outside_the_domain_raise(make) -> None: + """Loud, never a silently wrong answer: these are why `is_valid` has + nothing left to refuse.""" + x, gate_up_weight, down_weight = make() + with pytest.raises(RuntimeError): + with validating(op): + op(x, gate_up_weight, down_weight) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_draft_hidden_half.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_draft_hidden_half.py new file mode 100644 index 000000000000..94966647d7c4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_draft_hidden_half.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm_moe_dsa tp16 draft_hidden_half entry.""" + +import os +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.draft_hidden_half import ( # noqa: E501 + draft_hidden_half as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the glm_moe_dsa tp16 entries are certified on sm_107" +) + + +def _build(spec: dict, seed: int) -> dict: + """The op's arguments for one cell. + + The hidden state has the scale of a residual stream (standard deviation 8); the norm weight + sits around 1; the weight is scaled to unit-variance outputs. `out` starts as NaN, so an + output the kernel never stores fails the comparison. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + rows, k, n = spec["rows"], spec["k"], spec["n"] + x = (8.0 * torch.randn(rows, k, generator=g, device="cuda")).to(torch.bfloat16) + norm_weight = (1.0 + 0.1 * torch.randn(k, generator=g, device="cuda")).to(torch.bfloat16) + weight = (torch.randn(n, k, generator=g, device="cuda") * k**-0.5).to(torch.bfloat16) + in_kernel_norm = spec["in_kernel_norm"] + if not in_kernel_norm: + # The fallback's input: the row(s) already normalized and rounded to BF16. + xf = x.float() + x = (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + spec["eps"])).to(torch.bfloat16) + return dict( + x=x, + weight=weight, + out=torch.full((rows, n), float("nan"), dtype=torch.float32, device="cuda"), + norm_weight=norm_weight if in_kernel_norm else None, + eps=spec["eps"] if in_kernel_norm else None, + ) + + +def _compare(out: torch.Tensor, ref: torch.Tensor, in_kernel_norm: bool) -> None: + """FP32 outputs; the band depends on what can differ. + + Plain GEMM: the same BF16 operands, products exact in FP32, sums in another order (split-K + 8 over a 6144-long K against torch): about 1e-6 of a row's scale. Allowed: 2**-13 of the row + maximum element-wise and 2**-15 relative RMS, two orders above that and orders below a lost + or doubled K slice (about a third of the values). + + In-kernel norm: the kernel normalizes with FlashInfer's summation order and approximate + rsqrt, the reference with torch's, so a normalized element sitting next to a BF16 rounding + midpoint can round the other way (one BF16 ulp; a few FP32 ulps of rstd make that about one + element in 3 * 10**4, so a row sees none or one). One such flip moves every output by at + most 2**-7 of that element's product with its weight, about 2**-10.5 of the row maximum + with these operands. Allowed: 2**-8 of the row maximum element-wise and 2**-10 relative + RMS. A wrong eps, a missing norm weight or an unnormalized row is off by whole values. + """ + if in_kernel_norm: + assert_within_ulp(out, ref, element_ulp=1.0, rms_ulp=0.25, ulp=2.0**-8) + else: + assert_within_ulp(out, ref, element_ulp=1.0, rms_ulp=0.25, ulp=2.0**-13) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + x_before = args["x"].clone() + assert op(**args) is args["out"], "the op writes in place and returns what it wrote" + assert torch.equal(args["x"], x_before) + _compare(args["out"], op.reference(**args), cell.spec["in_kernel_norm"]) + + +def _flashinfer_norm_is_cute_dsl() -> bool: + """Whether trtllm::flashinfer_rmsnorm runs FlashInfer's CuTe DSL kernel, the one the in-kernel + norm reproduces. FlashInfer picks its CUDA kernel under FLASHINFER_USE_CUDA_NORM=1 and where + its own device check says so; without that check the default, CuTe DSL, is assumed.""" + if os.environ.get("FLASHINFER_USE_CUDA_NORM", "0") == "1": + return False + try: + from flashinfer.norm import _use_cuda_norm + except ImportError: + return True + return not _use_cuda_norm(torch.device("cuda", torch.cuda.current_device())) + + +def test_in_kernel_norm_is_flashinfer_rmsnorm_bit_for_bit() -> None: + """The point of the in-kernel norm: the result of FlashInfer's RMSNorm followed by this GEMM, + bit for bit, from one kernel.""" + if not _flashinfer_norm_is_cute_dsl(): + pytest.skip("FlashInfer runs its CUDA RMSNorm kernel here, not the one reproduced") + args = _build(op.CELLS[0].spec, seed=71) + with validating(op): + fused = op(**args) + normed = torch.ops.trtllm.flashinfer_rmsnorm(args["x"], args["norm_weight"], args["eps"]) + plain = op(normed, args["weight"], torch.empty_like(args["out"]), None, None) + assert torch.equal(fused, plain) + + +def test_cuda_graph_capture_and_replay() -> None: + """The draft loop captures the op on the auxiliary stream: it must replay on new contents of + the same buffers. Warm-up compiles the kernel outside the capture.""" + spec = op.CELLS[0].spec + args = _build(spec, seed=72) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + op(**args) + torch.cuda.current_stream().wait_stream(warmup) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(**args) + + args["x"].copy_(_build(spec, seed=73)["x"]) + args["out"].fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + _compare(args["out"], op.reference(**args), in_kernel_norm=True) + + +def test_guard_refuses_eps_without_the_norm_weight() -> None: + """eps alone would be ignored: the raw row would go into the GEMM.""" + args = _build(op.CELLS[0].spec, seed=74) + args["norm_weight"] = None + with pytest.raises(AssertionError, match="pass norm_weight and eps together"): + with validating(op): + op(**args) + + +def test_guard_refuses_an_out_over_the_weight() -> None: + """The owner CTAs store while other clusters still stream the weight: silently wrong.""" + args = _build(op.CELLS[0].spec, seed=75) + rows, n = args["out"].shape + args["out"] = args["weight"].view(torch.float32)[:rows, :n] + with pytest.raises(AssertionError, match="out must not share memory with weight"): + with validating(op): + op(**args) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_fused_a_gemm.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_fused_a_gemm.py new file mode 100644 index 000000000000..eed24372acbe --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_fused_a_gemm.py @@ -0,0 +1,239 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM fused_a_gemm entry. + +The entry owns the contract. This file turns a cell into tensors and holds the claims no cell +carries: the 2624 columns of the 2784-row call, determinism and the list as a pure hint, +agreement with the shared fused-A kernel, CUDA graph replay, the op's own rejections, and the +guards. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.fused_a_gemm import ( # noqa: E501 + HIDDEN_SIZE, + INDEXER_OUT_FEATURES, + OUT_FEATURES, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.fused_a_gemm import ( # noqa: E501 + fused_a_gemm as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402,F401 — needs the path declared above; fixture by name + assert_within_ulp, + certifying, + validating, +) + +assert torch.cuda.is_available(), "fused_a_gemm requires a CUDA device" + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM fused-A GEMM is certified on sm_107 only" +) + +DEV = "cuda" +BF16 = torch.bfloat16 + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """One GEMM's worth of accumulation-order difference, then one BF16 rounding on each side. + + Kernel and reference sum the same products (a BF16 x BF16 product is exact in FP32) in FP32 + and round once to BF16; only the summation order differs. With randn operands the outputs sit + near sqrt(6144) ~ 80, the two FP32 sums differ by about 1e-5 there (FP32 rounding over a + few dozen additions of partials up to that size), and the BF16 spacing there is 0.5. So an + output either rounds to the same BF16 value or lands one rounding step away, and a step is + at most 2 ulp (2^-8) of the row's largest magnitude. + + So: 4 ulp element-wise against the row scale (twice the worst rounding step), 2 ulp of + relative RMS (a step on 1 % of the outputs would read as about 0.15). A wrong k stage, row + or token is off by a fraction of the row scale and misses both by orders. + """ + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0) + + +def _tensor(shape, dtype, generator): + if dtype.is_floating_point: + return torch.randn(shape, generator=generator, device=DEV).to(dtype) + return torch.randint(0, 1 << 20, shape, generator=generator, device=DEV, dtype=dtype) + + +def _build(spec, seed): + """Materialize one cell: randn activations and weight, and the list tensors at their shapes.""" + g = torch.Generator(device=DEV).manual_seed(seed) + hidden_states = _tensor((spec["tokens"], HIDDEN_SIZE), BF16, g) + weight = _tensor((spec["out_features"], HIDDEN_SIZE), BF16, g) + l2_prefetch = [_tensor(shape, dtype, g) for shape, dtype in spec["l2_prefetch"]] + return hidden_states, weight, l2_prefetch + + +# ── the certified cells ─────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture + certifying(op) + hidden_states, weight, l2_prefetch = _build( + cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31 + ) + snapshots = [t.clone() for t in l2_prefetch] + out = op(hidden_states, weight, l2_prefetch) + assert out.shape == (hidden_states.shape[0], weight.shape[0]) + assert out.dtype == BF16 and out.is_contiguous() + _compare(out, op.reference(hidden_states, weight, l2_prefetch)) + # The list is a cache hint: read, never written. + for tensor, snapshot in zip(l2_prefetch, snapshots): + assert torch.equal(tensor, snapshot) + + +# ── the claims no cell carries ──────────────────────────────────────────── + + +@pytest.mark.parametrize("tokens", [1, 8]) +def test_indexer_rows_leave_the_kv_a_columns_bitwise_unchanged(tokens) -> None: + """A target keeps one 2784-row tensor and reads kv_a_proj_with_mqa as its first 2624 rows; + every 16-row tile is computed alike in both grids.""" + g = torch.Generator(device=DEV).manual_seed(tokens) + hidden_states = _tensor((tokens, HIDDEN_SIZE), BF16, g) + extended = _tensor((INDEXER_OUT_FEATURES, HIDDEN_SIZE), BF16, g) + kv_a = extended[:OUT_FEATURES] + assert kv_a.is_contiguous() and kv_a.data_ptr() == extended.data_ptr() + out_extended = op(hidden_states, extended, []) + out_kv_a = op(hidden_states, kv_a, []) + assert torch.equal(out_extended[:, :OUT_FEATURES], out_kv_a) + + +def test_deterministic_and_the_list_does_not_enter_the_result() -> None: + g = torch.Generator(device=DEV).manual_seed(11) + hidden_states = _tensor((8, HIDDEN_SIZE), BF16, g) + weight = _tensor((INDEXER_OUT_FEATURES, HIDDEN_SIZE), BF16, g) + l2_prefetch = [_tensor((4096, 2048), BF16, g), _tensor((6144, 1024), BF16, g)] + first = op(hidden_states, weight, l2_prefetch) + second = op(hidden_states, weight, l2_prefetch) + without_list = op(hidden_states, weight, []) + assert torch.equal(first, second) + assert torch.equal(first, without_list) + assert first.data_ptr() != second.data_ptr() + + +@pytest.mark.parametrize("tokens", [1, 8]) +def test_agrees_with_the_shared_fused_a_kernel(tokens) -> None: + """dsv3_fused_a_gemm_op (32-row tiles, 4 MMA warps) sums the same products in another order: + nearly every output is bitwise equal, the rest one BF16 rounding step apart.""" + g = torch.Generator(device=DEV).manual_seed(100 + tokens) + hidden_states = _tensor((tokens, HIDDEN_SIZE), BF16, g) + weight = _tensor((OUT_FEATURES, HIDDEN_SIZE), BF16, g) + out = op(hidden_states, weight, []) + shared = torch.ops.trtllm.dsv3_fused_a_gemm_op(hidden_states, weight.t(), None, None) + equal = (out == shared).float().mean().item() + print(f" {equal:.4%} of the outputs bitwise equal to dsv3_fused_a_gemm_op") + assert equal >= 0.99 + assert_within_ulp(out, shared, element_ulp=4.0, rms_ulp=2.0) + + +def test_cuda_graph_replay() -> None: + """The decode step runs as one CUDA graph: the tensor maps and the list are fixed at capture, + the contents of the captured activations are not.""" + g = torch.Generator(device=DEV).manual_seed(3) + hidden_states = _tensor((8, HIDDEN_SIZE), BF16, g) + weight = _tensor((INDEXER_OUT_FEATURES, HIDDEN_SIZE), BF16, g) + l2_prefetch = [_tensor((4096, 2048), BF16, g), _tensor((4, 256, 512), BF16, g)] + op(hidden_states, weight, l2_prefetch) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(hidden_states, weight, l2_prefetch) + for _ in range(2): + hidden_states.copy_(_tensor((8, HIDDEN_SIZE), BF16, g)) + graph.replay() + torch.cuda.synchronize() + _compare(out, op.reference(hidden_states, weight, l2_prefetch)) + + +# ── the op's own rejections ─────────────────────────────────────────────── + + +def test_op_rejects_what_the_kernel_cannot_run() -> None: + """Loud on the op's own, with no guard interposed.""" + g = torch.Generator(device=DEV).manual_seed(5) + hidden_states = _tensor((8, HIDDEN_SIZE), BF16, g) + weight = _tensor((OUT_FEATURES, HIDDEN_SIZE), BF16, g) + hidden_7168, weight_7168 = _tensor((8, 7168), BF16, g), _tensor((2112, 7168), BF16, g) + strided_hidden_states = _tensor((HIDDEN_SIZE, 8), BF16, g).t() + strided_weight = _tensor((HIDDEN_SIZE, OUT_FEATURES), BF16, g).t() + flat = torch.zeros(8 * HIDDEN_SIZE + 8, dtype=BF16, device=DEV) + misaligned = flat[1 : 1 + 8 * HIDDEN_SIZE].view(8, HIDDEN_SIZE) + cases = [ + ("fp16 hidden_states", lambda: op(hidden_states.half(), weight, [])), + ("fp32 weight", lambda: op(hidden_states, weight.float(), [])), + ("3D hidden_states", lambda: op(hidden_states.unsqueeze(0), weight, [])), + ("hidden 7168", lambda: op(hidden_7168, weight_7168, [])), + ("2640 rows", lambda: op(hidden_states, _tensor((2640, HIDDEN_SIZE), BF16, g), [])), + ("9 tokens", lambda: op(_tensor((9, HIDDEN_SIZE), BF16, g), weight, [])), + ("0 tokens", lambda: op(hidden_states[:0], weight, [])), + ("strided hidden_states", lambda: op(strided_hidden_states, weight, [])), + ("strided weight", lambda: op(hidden_states, strided_weight, [])), + ("misaligned hidden_states", lambda: op(misaligned, weight, [])), + ("cpu weight", lambda: op(hidden_states, weight.cpu(), [])), + ("five list tensors", lambda: op(hidden_states, weight, [weight] * 5)), + ("cpu list tensor", lambda: op(hidden_states, weight, [weight.cpu()])), + ("strided list tensor", lambda: op(hidden_states, weight, [weight.t()])), + ("misaligned list tensor", lambda: op(hidden_states, weight, [misaligned])), + ] + for tag, call in cases: + with pytest.raises((RuntimeError, NotImplementedError)): + call() + pytest.fail(f"{tag} was accepted") + + +# ── the guards ──────────────────────────────────────────────────────────── + + +def _ok(out_features=INDEXER_OUT_FEATURES, tokens=8): + g = torch.Generator(device=DEV).manual_seed(out_features + tokens) + return _tensor((tokens, HIDDEN_SIZE), BF16, g), _tensor((out_features, HIDDEN_SIZE), BF16, g) + + +def test_guard_refuses_another_hidden_size() -> None: + g = torch.Generator(device=DEV).manual_seed(1) + hidden_states, weight = _tensor((8, 7168), BF16, g), _tensor((2112, 7168), BF16, g) + with pytest.raises(AssertionError, match=r"\[tokens, 6144\]"): + with validating(op): + op(hidden_states, weight, []) + + +def test_guard_refuses_a_weight_without_an_instance() -> None: + hidden_states, _ = _ok() + g = torch.Generator(device=DEV).manual_seed(2) + with pytest.raises(AssertionError, match="2624 or 2784 rows"): + with validating(op): + op(hidden_states, _tensor((2640, HIDDEN_SIZE), BF16, g), []) + + +@pytest.mark.parametrize("tokens", [0, 9, 16]) +def test_guard_refuses_a_token_count_outside_one_tile(tokens) -> None: + hidden_states, weight = _ok(tokens=max(tokens, 1)) + hidden_states = hidden_states[:tokens] + with pytest.raises(AssertionError, match="1 to 8 tokens"): + with validating(op): + op(hidden_states, weight, []) + + +@pytest.mark.parametrize("out_features", [OUT_FEATURES, INDEXER_OUT_FEATURES]) +def test_guard_refuses_fewer_sms_than_ctas(out_features, monkeypatch) -> None: + """The op runs anyway, in two waves; the caller takes the shared op there.""" + hidden_states, weight = _ok(out_features) + ctas = out_features // 16 + + class _Props: + multi_processor_count = ctas - 1 + + monkeypatch.setattr(torch.cuda, "get_device_properties", lambda device=None: _Props()) + with pytest.raises(AssertionError, match=f"{ctas} CTAs need as many SMs"): + with validating(op): + op(hidden_states, weight, []) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_indexer_q_tail.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_indexer_q_tail.py new file mode 100644 index 000000000000..03675a813a89 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_indexer_q_tail.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm_moe_dsa tp16 indexer_q_tail entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.indexer_q_tail import ( # noqa: E501 + indexer_q_tail as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the glm_moe_dsa tp16 entries are certified on sm_107" +) + +#: GLM-5.2's fused-A output row with the indexer's wk (128) and weights_proj (32) rows: the +#: MLA's 2624 columns, then the indexer K, then the head weights. +_FUSED_A_WIDTH = 2784 +_WEIGHTS_COLUMN = 2624 + 128 +#: Rows of the test's RoPE table, and GLM-5.2's rope_theta. +_POSITIONS = 8192 +_THETA = 8.0e6 +#: An E4M3 NaN: what every Q byte holds until the kernel stores it. +_E4M3_NAN = 0x7F + + +def _indexer_cos_sin(positions: int, rope_dim: int) -> torch.Tensor: + """The indexer's RoPE table: per position the rope_dim / 2 cosines, then the sines.""" + inv_freq = _THETA ** (-torch.arange(0, rope_dim, 2, dtype=torch.float64) / rope_dim) + angle = torch.arange(positions, dtype=torch.float64)[:, None] * inv_freq[None, :] + return torch.cat((angle.cos(), angle.sin()), dim=-1).float().cuda().contiguous() + + +def _build(spec: dict, seed: int) -> dict: + """The op's arguments for one cell. + + q_a and the head weights are views of one fused-A output, as the target passes them. The + weight is scaled so that wq_b's outputs have a standard deviation of 2. The outputs start as + NaN, so anything the kernel does not store fails the comparison. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + m, k = spec["m"], spec["k"] + heads, head_dim = spec["n_heads"], spec["head_dim"] + fused_a = torch.randn(m, _FUSED_A_WIDTH, generator=g, device="cuda").to(torch.bfloat16) + weight = (torch.randn(heads * head_dim, k, generator=g, device="cuda") * 2.0 * k**-0.5).to( + torch.bfloat16 + ) + positions = torch.randint(0, _POSITIONS, (m,), generator=g, device="cuda", dtype=torch.int32) + q_fp8 = torch.full((m, heads, head_dim), _E4M3_NAN, dtype=torch.uint8, device="cuda") + return dict( + q_a=fused_a[:, :k], + weight=weight, + positions=positions, + cos_sin=_indexer_cos_sin(_POSITIONS, spec["rope_dim"]), + weights_in=fused_a[:, _WEIGHTS_COLUMN : _WEIGHTS_COLUMN + heads], + q_fp8=q_fp8.view(torch.float8_e4m3fn), + q_scale=torch.full((m, heads, 1), float("nan"), dtype=torch.float32, device="cuda"), + weights_out=torch.full((m, heads), float("nan"), dtype=torch.float32, device="cuda"), + eps=spec["eps"], + n_heads=heads, + head_dim=head_dim, + rope_dim=spec["rope_dim"], + interleave=spec["interleave"], + ue8m0=spec["ue8m0"], + weight_scale=spec["weight_scale"], + ) + + +def _compare(got: tuple, want: tuple) -> None: + """Scales and weights exact where the head's scale agrees; Q within one E4M3 ulp. + + wq_b's FP32 outputs differ from the reference's by the split-K summation order of a + 2048-long dot product and the approximate rsqrt (a few times 1e-6 relative), which changes + the BF16 rounding of about one value in a thousand by one BF16 ulp. Those flips reach the + E4M3 Q only where they cross an E4M3 rounding midpoint (1 in 16 of them) and reach a head's + UE8M0 scale only where the head's amax crosses a power of two (far rarer). Allowed: at most + 2 (token, head) scales differing, exact scaled weights where the scale agrees, every Q value + within one E4M3 ulp of the coarser of the two scales, at most max(8, 0.1 %) Q bytes + differing. A wrong position, pair, sign, scale rule or amax moves whole heads. + """ + q, scale, weights = got + q_ref, scale_ref, weights_ref = want + assert torch.isfinite(scale).all() and torch.isfinite(weights).all(), "an output is unwritten" + q_value = q.float() + assert torch.isfinite(q_value).all(), "a Q byte is unwritten (NaN)" + same_scale = scale == scale_ref + assert int((~same_scale).sum()) <= 2, f"{int((~same_scale).sum())} head scales differ" + same = same_scale.squeeze(-1) + assert torch.equal(weights[same], weights_ref[same]), "the scaled weights differ" + coarser = torch.maximum(scale, scale_ref) + got_value = q_value * scale + want_value = q_ref.float() * scale_ref + magnitude = (torch.maximum(got_value.abs(), want_value.abs()) / coarser).clamp_min(2.0**-6) + ulp = torch.exp2(torch.floor(torch.log2(magnitude)) - 3) * coarser + worst = ((got_value - want_value).abs() / ulp).max().item() + assert worst <= 1.0, f"a Q value is {worst} E4M3 ulp off" + differing = int((q.view(torch.uint8) != q_ref.view(torch.uint8)).sum()) + assert differing <= max(8, q.numel() // 1000), f"{differing} of {q.numel()} Q bytes differ" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + q_a_before, weights_in_before = args["q_a"].clone(), args["weights_in"].clone() + got = op(**args) + assert got[0] is args["q_fp8"] and got[1] is args["q_scale"] and got[2] is args["weights_out"] + assert torch.equal(args["q_a"], q_a_before) + assert torch.equal(args["weights_in"], weights_in_before) + _compare(got, op.reference(**args)) + + +def test_scales_are_powers_of_two() -> None: + """UE8M0: every scale is a power of two at or above amax / 448, so the Q never saturates.""" + args = _build(op.CELLS[0].spec, seed=51) + with validating(op): + q, scale, _ = op(**args) + mantissa = scale.view(torch.int32) & 0x007FFFFF + assert int(mantissa.abs().sum()) == 0, "a scale is not a power of two" + assert float(q.float().abs().max()) <= 448.0 + + +def test_guard_refuses_a_position_past_the_table() -> None: + args = _build(op.CELLS[0].spec, seed=52) + args["positions"][0] = _POSITIONS + with pytest.raises(AssertionError, match="every position must index a row"): + with validating(op): + op(**args) + + +def test_guard_refuses_a_q_over_q_a() -> None: + """The owner CTAs store while other clusters still read q_a: silently wrong.""" + args = _build(op.CELLS[0].spec, seed=53) + m, heads, head_dim = args["q_fp8"].shape + over = args["q_a"].view(torch.uint8)[:, : heads * head_dim].view(m, heads, head_dim) + args["q_fp8"] = over.view(torch.float8_e4m3fn) + with pytest.raises(AssertionError, match="q_fp8 must not share memory with q_a"): + with validating(op): + op(**args) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_low_m_gemm.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_low_m_gemm.py new file mode 100644 index 000000000000..ffe9d1b60579 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_low_m_gemm.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm_moe_dsa tp16 low_m_gemm entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.low_m_gemm import ( # noqa: E501 + low_m_gemm as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the glm_moe_dsa tp16 entries are certified on sm_107" +) + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """One GEMM's worth of accumulation-order difference. + + Kernel and reference multiply the same BF16 operands exactly in FP32 and differ only in the + order of the FP32 sums: the kernel's per-slice tensor-core accumulation and its fixed DSMEM + reduction of the K slices against torch's FP32 matmul. That is about 1e-6 relative to a + row's scale, which moves the one BF16 rounding of an output only where the FP32 sum sits + next to a rounding midpoint: rare one-ulp flips. One BF16 ulp of the row maximum is 2 units + of the band's 2**-8, so 4 units element-wise admits any such flip with margin, and 2 units of + relative RMS sits far above what a handful of flips adds. A lost or doubled K slice, a + wrong tile or a missing tail is off by the size of the values themselves, hundreds of units. + """ + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0) + + +def _build(spec: dict, seed: int): + """Operands of one cell: unit-variance activations, weights scaled to unit-variance outputs. + + `out` starts as NaN, so an output the kernel never stores fails the comparison. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + m, k, n = spec["m"], spec["k"], spec["n"] + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + weight = (torch.randn(n, k, generator=g, device="cuda") * k**-0.5).to(torch.bfloat16) + out = torch.full((m, n), float("nan"), dtype=torch.bfloat16, device="cuda") + return x, weight, out + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + x, weight, out = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + x_before, weight_before = x.clone(), weight.clone() + # Identity, not equality: the op writes the caller's buffer and returns it. + assert op(x, weight, out) is out, "the op writes in place and returns what it wrote" + assert torch.equal(x, x_before) and torch.equal(weight, weight_before) + _compare(out, op.reference(x, weight, out)) + + +def test_cuda_graph_capture_and_replay() -> None: + """The target captures the whole decode step: the op must replay on new contents of the + same buffers. Warm-up compiles the kernel outside the capture.""" + spec = op.CELLS[0].spec + x, weight, out = _build(spec, seed=31) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + op(x, weight, out) + torch.cuda.current_stream().wait_stream(warmup) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(x, weight, out) + + fresh, _, _ = _build(spec, seed=32) + x.copy_(fresh) + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + _compare(out, op.reference(x, weight, out)) + + +def _overlapping(spec: dict, operand: str): + """A cell's operands with `out` placed over `operand`'s memory.""" + x, weight, _ = _build(spec, seed=33) + m, n = spec["m"], spec["n"] + source = x if operand == "x" else weight + buffer = torch.empty(source.numel() + m * n, dtype=torch.bfloat16, device="cuda") + buffer[: source.numel()].copy_(source.view(-1)) + placed = buffer[: source.numel()].view(source.shape) + out = buffer[: m * n].view(m, n) + return (placed, weight, out) if operand == "x" else (x, placed, out) + + +@pytest.mark.parametrize("operand", ["x", "weight"]) +def test_guard_refuses_an_out_over_an_input(operand) -> None: + """The owner CTAs store while other clusters still read the inputs: silently wrong.""" + x, weight, out = _overlapping(op.CELLS[0].spec, operand) + with pytest.raises(AssertionError, match=f"out must not share memory with {operand}"): + with validating(op): + op(x, weight, out) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_o_proj_publish.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_o_proj_publish.py new file mode 100644 index 000000000000..6a856d09fe8e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_o_proj_publish.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm_moe_dsa tp16 o_proj_publish entry. + +One GPU stands in for the group: `lamport` is a plain device buffer laid out like an MNNVL +one-shot workspace (three Lamport buffers, every element the idle -0.0) and `flags` a plain int32 +tensor, so the test can read back exactly which slots the op stored and that it stored nothing +else. On a real group `lamport` is the multicast address and the same stores reach every rank. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.low_m_gemm import ( # noqa: E501 + low_m_gemm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.o_proj_publish import ( # noqa: E501 + o_proj_publish as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the glm_moe_dsa tp16 entries are certified on sm_107" +) + +#: Lamport buffers of a one-shot workspace. +_BUFFERS = 3 +#: Bytes past the slots in each buffer: a store past the end of the slot block lands there. +_SLACK_BYTES = 4096 +#: The BF16 bit pattern of -0.0, the idle value of a Lamport buffer, as int16. +_IDLE = -32768 + + +def _build(spec: dict, seed: int) -> dict: + """The op's arguments for one cell: unit-variance attention outputs, a weight scaled to + unit-variance outputs, a workspace with the slots of every token and rank plus slack.""" + g = torch.Generator(device="cuda").manual_seed(seed) + m, k, n, world = spec["m"], spec["k"], spec["n"], spec["world_size"] + bytes_per_buffer = m * world * n * 2 + _SLACK_BYTES + current = spec["current"] + flags = torch.zeros(16, dtype=torch.int32, device="cuda") + flags[:4] = torch.tensor([current, (current + 2) % _BUFFERS, bytes_per_buffer, 1]) + return dict( + x=torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16), + weight=(torch.randn(n, k, generator=g, device="cuda") * k**-0.5).to(torch.bfloat16), + out=torch.empty(m, n, dtype=torch.bfloat16, device="cuda"), + lamport=torch.full( + (_BUFFERS * bytes_per_buffer // 2,), -0.0, dtype=torch.bfloat16, device="cuda" + ), + flags=flags, + world_size=world, + rank=spec["rank"], + ) + + +def _slots(args: dict) -> tuple[torch.Tensor, torch.Tensor]: + """(this rank's [M, N] slots in the current buffer, a mask of every element of `lamport` + outside them).""" + m, n = args["x"].shape[0], args["weight"].shape[0] + world, rank = args["world_size"], args["rank"] + current, _, bytes_per_buffer, _ = args["flags"][:4].tolist() + buffer = args["lamport"].view(_BUFFERS, bytes_per_buffer // 2)[current] + slots = buffer[: m * world * n].view(m, world, n) + elsewhere = torch.ones(args["lamport"].numel(), dtype=torch.bool, device="cuda") + mask = elsewhere.view(_BUFFERS, -1)[current][: m * world * n].view(m, world, n) + mask[:, rank, :] = False + return slots[:, rank, :], elsewhere + + +def _compare(published: torch.Tensor, ref: torch.Tensor) -> None: + """One GEMM's worth of accumulation-order difference, as for low_m_gemm. + + The FP32 sums of the kernel (split-K 2 over a 1024-long K) and of torch differ in order by + about 1e-6 relative, which flips the BF16 rounding of a rare output by one ulp: 2 units of + the band's 2**-8 at the row maximum, so 4 units element-wise and 2 units of relative RMS. A + value in the wrong slot, a lost K slice or a wrong rounding is off by whole values. + """ + assert_within_ulp(published, ref, element_ulp=4.0, rms_ulp=2.0) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + assert op(**args) is None + published, elsewhere = _slots(args) + assert bool((published.view(torch.int16) != _IDLE).all()), "a slot holds -0.0 (unwritten)" + idle = args["lamport"].view(torch.int16)[elsewhere] + assert bool((idle == _IDLE).all()), "the op stored outside this rank's slots" + _compare(published, op.reference(**args)) + + +def test_published_values_are_the_gemm_output_bit_for_bit() -> None: + """The publish epilogue rounds as the plain epilogue does: the slots hold low_m_gemm's + output at the same tactic, -0.0 written as +0.0.""" + args = _build(op.CELLS[0].spec, seed=61) + with validating(op): + op(**args) + published, _ = _slots(args) + plain = low_m_gemm(args["x"], args["weight"], torch.empty_like(args["out"])) + plain = torch.where(plain == 0, torch.zeros_like(plain), plain) + assert torch.equal(published.view(torch.int16), plain.view(torch.int16)) + + +def test_cuda_graph_capture_and_replay() -> None: + """Captured once, the op publishes into whichever buffer the flags mark at replay.""" + args = _build(op.CELLS[0].spec, seed=62) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + op(**args) + torch.cuda.current_stream().wait_stream(warmup) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(**args) + + args["lamport"].fill_(-0.0) + args["flags"][0] = (int(args["flags"][0]) + 1) % _BUFFERS + args["x"].copy_(_build(op.CELLS[0].spec, seed=63)["x"]) + graph.replay() + torch.cuda.synchronize() + published, elsewhere = _slots(args) + assert bool((args["lamport"].view(torch.int16)[elsewhere] == _IDLE).all()) + _compare(published, op.reference(**args)) + + +def test_guard_refuses_a_rank_outside_the_group() -> None: + """Rank 16 of 16 lands in the next token's slots.""" + args = _build(op.CELLS[0].spec, seed=64) + args["rank"] = args["world_size"] + with pytest.raises(AssertionError, match="outside a group of"): + with validating(op): + op(**args) + + +def test_guard_refuses_a_buffer_smaller_than_the_slots() -> None: + """The slots past the buffer's end land in the next Lamport buffer.""" + args = _build(op.CELLS[0].spec, seed=65) + m, n = args["x"].shape[0], args["weight"].shape[0] + args["flags"][2] = m * args["world_size"] * n * 2 - 16 + with pytest.raises(AssertionError, match="smaller than the"): + with validating(op): + op(**args) + + +def test_guard_refuses_a_buffer_index_out_of_range() -> None: + args = _build(op.CELLS[0].spec, seed=66) + args["flags"][0] = _BUFFERS + with pytest.raises(AssertionError, match="current buffer index"): + with validating(op): + op(**args) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_q_b_fp8_rope.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_q_b_fp8_rope.py new file mode 100644 index 000000000000..cd70e5d27797 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_q_b_fp8_rope.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm_moe_dsa tp16 q_b_fp8_rope entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.q_b_fp8_rope import ( # noqa: E501 + q_b_fp8_rope as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the glm_moe_dsa tp16 entries are certified on sm_107" +) + +#: GLM-5.2's fused-A output row: q_lora_rank 2048, kv_lora_rank 512, qk_rope_head_dim 64. +_FUSED_A_WIDTH = 2624 +#: Rows of the test's RoPE table, and GLM-5.2's rope_theta. +_POSITIONS = 8192 +_THETA = 8.0e6 +#: An E4M3 NaN: what every output byte holds until the kernel stores it. +_E4M3_NAN = 0x7F + + +def _mla_cos_sin(positions: int, rope_dim: int) -> torch.Tensor: + """The MLA RoPE table in the layout the attention backend builds for MLA: per position + rope_dim (cos, sin) pairs, the rope_dim / 2 frequencies theta ** (-2i / rope_dim) twice.""" + inv_freq = _THETA ** (-torch.arange(0, rope_dim, 2, dtype=torch.float64) / rope_dim) + angle = torch.arange(positions, dtype=torch.float64)[:, None] * inv_freq[None, :] + angle = torch.cat((angle, angle), dim=-1) + return torch.stack((angle.cos(), angle.sin()), dim=-1).float().cuda().contiguous() + + +def _build(spec: dict, seed: int) -> dict: + """The op's arguments for one cell. + + q_a is the q_lora_rank-column row view of a fused-A output, as the target passes it. The + weight is scaled so that the FP32 values before the conversion have a standard deviation of + 4 (times q_scale): E4M3 normals of every exponent from 2**-6 up, and subnormals below. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + m, k, seq_len = spec["m"], spec["k"], spec["seq_len"] + n = spec["heads"] * spec["head_dim"] + fused_a = torch.randn(m, _FUSED_A_WIDTH, generator=g, device="cuda").to(torch.bfloat16) + weight = (torch.randn(n, k, generator=g, device="cuda") * 4.0 * k**-0.5).to(torch.bfloat16) + kv_lens = torch.randint( + seq_len, + _POSITIONS + 1, + (m // seq_len,), + generator=g, + device="cuda", + dtype=torch.int32, + ) + return dict( + q_a=fused_a[:, :k], + weight=weight, + out=torch.full((m, n), _E4M3_NAN, dtype=torch.uint8, device="cuda"), + kv_lens=kv_lens, + cos_sin=_mla_cos_sin(_POSITIONS, spec["rope_dim"]), + q_scale=torch.full((1,), spec.get("q_scale", 1.0), dtype=torch.float32, device="cuda"), + eps=spec["eps"], + head_dim=spec["head_dim"], + rope_dim=spec["rope_dim"], + seq_len=seq_len, + plan=spec["plan"], + ) + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Every E4M3 value within one E4M3 ulp, and only the rare ones not equal. + + Before its single rounding, the kernel's FP32 value differs from the reference's by the + split-K summation order of a 2048-long dot product (about 3e-6 relative), the approximate + rsqrt of the row RMS (about 2**-22) and the contraction of the RoPE products into FMAs (one + FP32 ulp): a few times 1e-6 relative. E4M3 spaces its values 2**-4 to 2**-3 apart + (relative), so a value converts differently only within that distance of a rounding + midpoint: a probability of about 1e-4 per element, about two of the 18432 values of a verify + cell. Allowed: one E4M3 ulp per element, and at most max(4, 0.1 %) differing bytes. A wrong + position, pair, sign or scale moves most RoPE values by many ulps, and a lost K slice moves + every value. + """ + got = out.view(torch.float8_e4m3fn).float() + want = ref.view(torch.float8_e4m3fn).float() + assert torch.isfinite(got).all(), "the op left an output byte unwritten (NaN) or wrote NaN" + magnitude = torch.maximum(got.abs(), want.abs()).clamp_min(2.0**-6) + ulp = torch.exp2(torch.floor(torch.log2(magnitude)) - 3) + worst = ((got - want).abs() / ulp).max().item() + assert worst <= 1.0, f"an E4M3 value is {worst} ulp off" + differing = int((out != ref).sum()) + assert differing <= max(4, out.numel() // 1000), f"{differing} of {out.numel()} bytes differ" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + q_a_before, weight_before = args["q_a"].clone(), args["weight"].clone() + assert op(**args) is args["out"], "the op writes in place and returns what it wrote" + assert torch.equal(args["q_a"], q_a_before) and torch.equal(args["weight"], weight_before) + _compare(args["out"], op.reference(**args)) + + +def test_q_scale_is_applied_and_saturates() -> None: + """A scale of 32 pushes the values past 448: the conversion saturates instead of producing + NaN, and the scale multiplies before it (c1 runs at 1, where the scale is invisible).""" + spec = dict(op.CELLS[0].spec, q_scale=32.0) + args = _build(spec, seed=41) + with validating(op): + op(**args) + got = args["out"].view(torch.float8_e4m3fn).float() + assert (got.abs() == 448.0).any(), "no value reached the saturation bound" + _compare(args["out"], op.reference(**args)) + + +def test_plans_agree() -> None: + """The plans change the launch and the FP32 summation order, not the definition.""" + spec = op.CELLS[0].spec + results = {} + for plan in ("regular", "beside_topk", "early_gather"): + args = _build(dict(spec, plan=plan), seed=43) + op(**args) + results[plan] = args["out"] + _compare(results["beside_topk"], results["regular"]) + _compare(results["early_gather"], results["regular"]) + + +def test_cuda_graph_capture_and_replay() -> None: + """The target captures the whole decode step: the op must replay on new contents of the same + buffers. Warm-up compiles the kernel and builds its placeholders outside the capture.""" + spec = op.CELLS[0].spec + args = _build(spec, seed=44) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + op(**args) + torch.cuda.current_stream().wait_stream(warmup) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(**args) + + fresh = _build(spec, seed=45) + args["q_a"].copy_(fresh["q_a"]) + args["kv_lens"].copy_(fresh["kv_lens"]) + args["out"].fill_(_E4M3_NAN) + graph.replay() + torch.cuda.synchronize() + _compare(args["out"], op.reference(**args)) + + +def test_guard_refuses_a_table_of_another_row_width() -> None: + """A table with rope_dim / 2 pairs per position (no repeated half) is read at the wrong rows.""" + args = _build(op.CELLS[0].spec, seed=46) + args["cos_sin"] = args["cos_sin"][:, : args["rope_dim"] // 2].contiguous() + with pytest.raises(AssertionError, match="cos_sin must be"): + with validating(op): + op(**args) + + +def test_guard_refuses_a_position_past_the_table() -> None: + args = _build(op.CELLS[0].spec, seed=47) + args["kv_lens"].fill_(_POSITIONS + 1) + with pytest.raises(AssertionError, match="cos/sin table must hold every position"): + with validating(op): + op(**args) + + +def test_guard_refuses_kv_lens_shorter_than_the_step() -> None: + """kv_lens below seq_len gives negative positions: reads before the table.""" + args = _build(op.CELLS[0].spec, seed=48) + args["kv_lens"].fill_(args["seq_len"] - 1) + with pytest.raises(AssertionError, match="kv_lens must cover the step's own tokens"): + with validating(op): + op(**args) + + +def test_guard_refuses_an_out_over_q_a() -> None: + """The owner CTAs store while other clusters still read q_a: silently wrong.""" + args = _build(op.CELLS[0].spec, seed=49) + m, n = args["out"].shape + args["out"] = args["q_a"].view(torch.uint8)[:, :n] + assert args["out"].shape == (m, n) + with pytest.raises(AssertionError, match="out must not share memory with q_a"): + with validating(op): + op(**args) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_router_gemm.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_router_gemm.py new file mode 100644 index 000000000000..29a981ee05a2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_router_gemm.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM router_gemm entry. + +The entry owns the contract. This file turns a cell into tensors and holds the claims no cell +carries: bitwise agreement with the shared router GEMV, determinism, every token count up to the +ceiling, CUDA graph replay, the op's own rejections, and the guards. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.router_gemm import ( # noqa: E501 + HIDDEN_SIZE, + MAX_TOKENS, + NUM_EXPERTS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.router_gemm import ( # noqa: E501 + router_gemm as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402,F401 — needs the path declared above; fixture by name + ULP, + assert_within_ulp, + certifying, + validating, +) + +assert torch.cuda.is_available(), "router_gemm requires a CUDA device" + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM router GEMV is certified on sm_107 only" +) + +DEV = "cuda" +BF16 = torch.bfloat16 + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """FP32 accumulation against an FP64 reference, scored in FP32 ulps. + + Both sides sum the same products (a BF16 x BF16 product is exact in FP32); the reference in + FP64, the kernel in FP32: per thread a chain of 48 FMAs, then a 5-level warp butterfly and the + four warp sums. With randn operands the logits sit near sqrt(6144) ~ 80 (row maxima ~230), + and the FP32 roundings of those additions add up to about 2.5e-5 per logit: about 3 ulp + (2^-23) of relative RMS and about 3 ulp of the row maximum at worst. + + So: 64 ulp element-wise against the row scale and 32 ulp of relative RMS, tenfold over the + estimate. A BF16 rounding of the logits alone would be 2^-9, four orders above, and a missing + k chunk, expert or token misses by more. + """ + assert_within_ulp(out, ref, element_ulp=64.0, rms_ulp=32.0, ulp=ULP[torch.float32]) + + +def _build(tokens, seed): + """randn activations and gate weight, both BF16.""" + g = torch.Generator(device=DEV).manual_seed(seed) + hidden_states = torch.randn(tokens, HIDDEN_SIZE, generator=g, device=DEV).to(BF16) + weight = torch.randn(NUM_EXPERTS, HIDDEN_SIZE, generator=g, device=DEV).to(BF16) + return hidden_states, weight + + +# ── the certified cells ─────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture + certifying(op) + hidden_states, weight = _build(cell.spec["tokens"], zlib.crc32(cell.why.encode()) % 2**31) + logits = op(hidden_states, weight) + assert logits.shape == (hidden_states.shape[0], NUM_EXPERTS) + assert logits.dtype == torch.float32 and logits.is_contiguous() + _compare(logits, op.reference(hidden_states, weight)) + + +# ── the claims no cell carries ──────────────────────────────────────────── + + +@pytest.mark.parametrize("tokens", range(1, MAX_TOKENS + 1)) +def test_bitwise_the_shared_router_gemv(tokens) -> None: + """dsv3_router_gemm_op runs the same loads and FMA chain; only the trigger and the L2 + prefetch differ, and neither touches the arithmetic. Every token count is an instance.""" + hidden_states, weight = _build(tokens, seed=tokens) + logits = op(hidden_states, weight) + shared = torch.ops.trtllm.dsv3_router_gemm_op(hidden_states, weight.t(), None, torch.float32) + assert torch.equal(logits, shared) + _compare(logits, op.reference(hidden_states, weight)) + + +def test_deterministic() -> None: + hidden_states, weight = _build(8, seed=21) + first = op(hidden_states, weight) + second = op(hidden_states, weight) + assert torch.equal(first, second) and first.data_ptr() != second.data_ptr() + + +def test_cuda_graph_replay() -> None: + """The decode step runs as one CUDA graph: the captured activations change between replays.""" + hidden_states, weight = _build(1, seed=4) + op(hidden_states, weight) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + logits = op(hidden_states, weight) + for seed in (5, 6): + hidden_states.copy_(_build(1, seed)[0]) + graph.replay() + torch.cuda.synchronize() + _compare(logits, op.reference(hidden_states, weight)) + + +# ── the op's own rejections ─────────────────────────────────────────────── + + +def test_op_rejects_what_the_kernel_cannot_run() -> None: + """Loud on the op's own, with no guard interposed.""" + hidden_states, weight = _build(8, seed=9) + g = torch.Generator(device=DEV).manual_seed(10) + hidden_7168 = torch.randn(8, 7168, generator=g, device=DEV).to(BF16) + weight_7168 = torch.randn(NUM_EXPERTS, 7168, generator=g, device=DEV).to(BF16) + weight_128 = weight[:128] + hidden_17 = torch.randn(17, HIDDEN_SIZE, generator=g, device=DEV).to(BF16) + strided_hidden_states = torch.randn(HIDDEN_SIZE, 8, generator=g, device=DEV).to(BF16).t() + flat = torch.zeros(8 * HIDDEN_SIZE + 8, dtype=BF16, device=DEV) + misaligned = flat[1 : 1 + 8 * HIDDEN_SIZE].view(8, HIDDEN_SIZE) + cases = [ + ("fp16 hidden_states", lambda: op(hidden_states.half(), weight)), + ("fp32 weight", lambda: op(hidden_states, weight.float())), + ("3D hidden_states", lambda: op(hidden_states.unsqueeze(0), weight)), + ("hidden 7168", lambda: op(hidden_7168, weight_7168)), + ("128 experts", lambda: op(hidden_states, weight_128)), + ("17 tokens", lambda: op(hidden_17, weight)), + ("0 tokens", lambda: op(hidden_states[:0], weight)), + ("strided hidden_states", lambda: op(strided_hidden_states, weight)), + ("weight.t()", lambda: op(hidden_states, weight.t())), + ("misaligned hidden_states", lambda: op(misaligned, weight)), + ("cpu weight", lambda: op(hidden_states, weight.cpu())), + ] + for tag, call in cases: + with pytest.raises((RuntimeError, NotImplementedError)): + call() + pytest.fail(f"{tag} was accepted") + + +# ── the guards ──────────────────────────────────────────────────────────── + + +def test_guard_refuses_another_hidden_size() -> None: + g = torch.Generator(device=DEV).manual_seed(1) + hidden_states = torch.randn(1, 7168, generator=g, device=DEV).to(BF16) + weight = torch.randn(NUM_EXPERTS, 7168, generator=g, device=DEV).to(BF16) + with pytest.raises(AssertionError, match=r"\[tokens, 6144\]"): + with validating(op): + op(hidden_states, weight) + + +def test_guard_refuses_another_expert_count() -> None: + hidden_states, weight = _build(1, seed=2) + with pytest.raises(AssertionError, match=r"weight must be \[256, 6144\]"): + with validating(op): + op(hidden_states, weight[:128]) + + +@pytest.mark.parametrize("tokens", [0, 17, 32]) +def test_guard_refuses_a_token_count_without_an_instance(tokens) -> None: + hidden_states, weight = _build(max(tokens, 1), seed=3) + with pytest.raises(AssertionError, match="1 to 16 tokens"): + with validating(op): + op(hidden_states[:tokens], weight) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_lowm_bf16_moe.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_lowm_bf16_moe.py new file mode 100644 index 000000000000..9c609339c90c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_lowm_bf16_moe.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm_lowm_bf16_moe entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.moe.lowm_bf16_moe import ( # noqa: E501 + glm_lowm_bf16_moe as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402,F401 — needs the path declared above + assert_within_ulp, + certifying, + validating, +) + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" +) + +NUM_EXPERTS = 256 +TOP_K = 8 +ROUTED_SCALING_FACTOR = 2.5 + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """One accumulation order's worth of BF16 rounding flips. + + The reference rounds at the op's points (FC1, SwiGLU, each scaled expert output, the sum), so the two differ + only where the FP32 accumulation order of a dot product moves a BF16 rounding of an intermediate. + + Over 6144-term FC1 dots of unit-scale operands the order changes the FP32 sum by a few 1e-6 relative, which + flips the BF16 rounding of about one gate or up value in a thousand; a flipped activation moves its expert's + FC2 outputs by about one BF16 ulp of the activation times a 1/sqrt(128)-scale weight, far below one ulp of the + output. What reaches the output is then the occasional flip of a term's or the sum's own rounding: one ulp of + that element, which is at most one ulp of the row's scale. + + So: 2 ulp element-wise against the row scale, 0.5 ulp of relative RMS. A wrong expert or weight moves the + output by about its own size, hundreds of ulp. + """ + assert_within_ulp(out, ref, element_ulp=2.0, rms_ulp=0.5) + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + hidden, inter = spec["hidden"], spec["inter"] + x = torch.randn(1, hidden, generator=g, device="cuda", dtype=torch.bfloat16) + w3_w1 = torch.randn( + NUM_EXPERTS, 2 * inter, hidden, generator=g, device="cuda", dtype=torch.bfloat16 + ).mul_(hidden**-0.5) + w2 = torch.randn( + NUM_EXPERTS, hidden, inter, generator=g, device="cuda", dtype=torch.bfloat16 + ).mul_(inter**-0.5) + routing = spec["routing"] + if routing == "random": + logits = 2.0 * torch.randn(1, NUM_EXPERTS, generator=g, device="cuda") + bias = 0.1 * torch.randn(NUM_EXPERTS, generator=g, device="cuda") + elif routing == "ties": + # Half-integer logits in [-2, 2] and no bias: about 28 experts share the top key, so which 8 of them are + # selected, and in which order, is the tie-break. + logits = torch.randint(-4, 5, (1, NUM_EXPERTS), generator=g, device="cuda").float() / 2.0 + bias = torch.zeros(NUM_EXPERTS, device="cuda") + else: + assert routing == "tiny" + # sigmoid(-16 +- 1) is about 1e-7, so the 8 selected scores sum to about 1e-6 < 2^-13. + logits = -16.0 + 0.5 * torch.randn(1, NUM_EXPERTS, generator=g, device="cuda") + bias = 0.1 * torch.randn(NUM_EXPERTS, generator=g, device="cuda") + return x, logits.contiguous(), bias.float().contiguous(), w3_w1, w2 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + out = op(*args, TOP_K, ROUTED_SCALING_FACTOR) + torch.cuda.synchronize() + _compare(out, op.reference(*args, TOP_K, ROUTED_SCALING_FACTOR)) + + +def test_cuda_graph_replay_matches_eager() -> None: + """The draft loop runs inside the decode graph: a replay with new inputs gives the eager result bitwise.""" + spec = op.CELLS[0].spec + x, logits, bias, w3_w1, w2 = _build(spec, seed=1) + static_x, static_logits = x.clone(), logits.clone() + out = op(static_x, static_logits, bias, w3_w1, w2, TOP_K, ROUTED_SCALING_FACTOR) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_out = op(static_x, static_logits, bias, w3_w1, w2, TOP_K, ROUTED_SCALING_FACTOR) + x2, logits2, _, _, _ = _build(spec, seed=2) + static_x.copy_(x2) + static_logits.copy_(logits2) + graph.replay() + eager = op(x2, logits2, bias, w3_w1, w2, TOP_K, ROUTED_SCALING_FACTOR) + torch.cuda.synchronize() + assert torch.equal(graph_out, eager) + assert not torch.equal(out, eager), "the replay must have read the new inputs" + + +def _small(hidden=2048, inter=64, tokens=1): + """Operands of the right layout at a small supported shape. + + The weights stay uninitialized: only the refusals, which happen before any launch, are exercised. + """ + x = torch.zeros(tokens, hidden, dtype=torch.bfloat16, device="cuda") + logits = torch.zeros(tokens, NUM_EXPERTS, dtype=torch.float32, device="cuda") + bias = torch.zeros(NUM_EXPERTS, dtype=torch.float32, device="cuda") + w3_w1 = torch.empty(NUM_EXPERTS, 2 * inter, hidden, dtype=torch.bfloat16, device="cuda") + w2 = torch.empty(NUM_EXPERTS, hidden, inter, dtype=torch.bfloat16, device="cuda") + return x, logits, bias, w3_w1, w2 + + +def test_refuses_more_than_one_token() -> None: + """Draft step 0 without the row select has 8 rows: the caller must take another MoE path. + + The op says so rather than computing token 0 alone. + """ + with pytest.raises(RuntimeError, match="takes one token"): + with validating(op): + op(*_small(tokens=2), TOP_K, ROUTED_SCALING_FACTOR) + + +def test_refuses_an_expert_slice_it_does_not_serve() -> None: + """The FC2 kernel walks the expert slice in 64-element K steps.""" + with pytest.raises(RuntimeError, match="unsupported hidden size"): + with validating(op): + op(*_small(inter=96), TOP_K, ROUTED_SCALING_FACTOR) + + +def test_refuses_another_top_k() -> None: + with pytest.raises(RuntimeError, match="top-8"): + with validating(op): + op(*_small(), 6, ROUTED_SCALING_FACTOR) + + +def test_refuses_a_misaligned_token() -> None: + """The FC1 kernel reads x in 16-byte pieces.""" + _, logits, bias, w3_w1, w2 = _small() + storage = torch.zeros(1, 2048 + 1, dtype=torch.bfloat16, device="cuda") + misaligned = storage[:, 1:] + assert misaligned.is_contiguous() and misaligned.data_ptr() % 16 != 0 + with pytest.raises(RuntimeError, match="16-byte aligned"): + with validating(op): + op(misaligned, logits, bias, w3_w1, w2, TOP_K, ROUTED_SCALING_FACTOR) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_routed_experts.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_routed_experts.py new file mode 100644 index 000000000000..57a7f2653d88 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_routed_experts.py @@ -0,0 +1,646 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the GLM-5.2 tp16 routed_experts entry.""" + +import zlib + +import pytest +import torch +from torch.profiler import ProfilerActivity, profile + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.moe.routed_experts import ( # noqa: E501 + ACT_TYPE_SWIGLU, + FUSED_ROUTER_WORKSPACE_WORDS, + ROUTING_METHOD_DEEPSEEK_V3, + make_fused_router_workspace, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.moe.routed_experts import ( # noqa: E501 + routed_experts as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402,F401 — needs the path declared above + assert_within_ulp, + certifying, + validating, +) + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" +) + +DEV = "cuda" +HIDDEN, INTER, EXPERTS, TOP_K = 6144, 128, 256, 8 +ROUTED_SCALE = 2.5 +SV = 16 # NVFP4 block size: one e4m3 scale per 16 elements along K +E2M1_MAX, E4M3_MAX = 6.0, 448.0 +# Expert-selection scores (sigmoid + bias) of the 8th and 9th expert closer than +# this leave the reference's choice to the sigmoid's last bit; a case compared +# against the reference is drawn clear of it. +SELECTION_MARGIN = 1e-6 +# Long enough (tens of ms) that a kernel not waiting for the side stream reads +# the stale input long before the side stream writes the real one. +SLEEP_CYCLES = 1 << 27 +_BUFFER_WORDS = (FUSED_ROUTER_WORKSPACE_WORDS - 4) // 2 +_EPOCH_WORD = 2 * _BUFFER_WORDS + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """16 ulp per element against the row's own scale, 2 ulp relative RMS. + + The band of the catalog `fp4_block_scale_moe_runner` entry, which compares + the same kernels against the same reference over bit-identical NVFP4 + operands: they differ only in accumulation order and in whether a marginal + FC1 value lands on the same e2m1 code at the FC2-input requantization. Its + measured worst over every geometry it covers is 9.97 ulp element-wise and + 0.88 ulp RMS, 4.68 / 0.70 for DeepSeek-R1's 256-expert stack; a reference + that skips the requantization lands 23 ulp RMS out. GLM's FC2 sums over 128 + terms where R1's sums over 2048, so accumulation order moves it less. + Routing adds one source those comparisons did not have: a routing weight + the kernel rounds to bf16 one ulp away from the reference's moves one + slot's share of a row by 2^-8, under 1 ulp of the row. A different expert + set is hundreds of ulp out. + """ + assert_within_ulp(out, ref, element_ulp=16.0, rms_ulp=2.0) + + +# ── operands ────────────────────────────────────────────────────────────── + + +def _rand_nvfp4(e: int, n: int, k: int, gen: torch.Generator): + """Random NVFP4 expert stack: `(packed [E, N, K/2] uint8, codes [E, N, K] + uint8, scales [E, N, K/16] float8_e4m3fn)`. Codes and power-of-two scales + are drawn directly, so every weight's fp32 value is exact.""" + codes = torch.randint(0, 16, (e, n, k), dtype=torch.uint8, device=DEV, generator=gen) + exps = torch.randint(-2, 1, (e, n, k // SV), device=DEV, generator=gen, dtype=torch.int32) + scales = torch.exp2(exps.float()).to(torch.float8_e4m3fn) + packed = (codes[..., 0::2] | (codes[..., 1::2] << 4)).contiguous() + return packed, codes, scales + + +def _blk32_perm(m: int) -> torch.Tensor: + """Gather index of the 32-row block shuffle: within each block of 32 rows, + source row `4u + v` lands at destination row `8v + u`.""" + j = torch.arange(32) + idx = torch.empty(32, dtype=torch.long) + idx[(j % 4) * 8 + j // 4] = j + return (idx.repeat(m // 32) + torch.arange(m // 32).repeat_interleave(32) * 32).to(DEV) + + +def _interleave_perm(m: int) -> torch.Tensor: + """Gather index interleaving `[up | gate]` rows into up0, gate0, up1, ...""" + p = torch.empty(m, dtype=torch.long) + p[0::2] = torch.arange(0, m // 2) + p[1::2] = torch.arange(m // 2, m) + return p.to(DEV) + + +def _swizzle(s: torch.Tensor) -> torch.Tensor: + """Per-expert 128x4 block-scale swizzle of `[E, M, C]` uint8 scales.""" + e, m, c = s.shape + return s.reshape(e, m // 128, 4, 32, c // 4, 4).permute(0, 1, 4, 3, 2, 5).reshape(e, m, c) + + +def _prepare(packed: torch.Tensor, scales: torch.Tensor, perm: torch.Tensor): + """Shuffle the rows of a stack and its scales, then swizzle the scales.""" + w = torch.index_select(packed, 1, perm).contiguous() + s = torch.index_select(scales.view(torch.uint8), 1, perm) + return w, _swizzle(s).contiguous().view(torch.float8_e4m3fn) + + +_STACK = None + + +def _stack(): + """GLM-5.2's expert stack at MoE TP16 -- 256 experts, H 6144, I 128 -- + built once: the kernel operands and the reference's `(codes, scales)`.""" + global _STACK + if _STACK is None: + gen = torch.Generator(device=DEV).manual_seed(5216) + up_p, up_c, up_s = _rand_nvfp4(EXPERTS, INTER, HIDDEN, gen) + gt_p, gt_c, gt_s = _rand_nvfp4(EXPERTS, INTER, HIDDEN, gen) + dn_p, dn_c, dn_s = _rand_nvfp4(EXPERTS, HIDDEN, INTER, gen) + fc1_perm = _interleave_perm(2 * INTER)[_blk32_perm(2 * INTER)] + w1, s1 = _prepare(torch.cat([up_p, gt_p], dim=1), torch.cat([up_s, gt_s], dim=1), fc1_perm) + w2, s2 = _prepare(dn_p, dn_s, _blk32_perm(HIDDEN)) + _STACK = dict( + gemm1_weights=w1, + gemm1_weights_scale=s1, + gemm2_weights=w2, + gemm2_weights_scale=s2, + up=(up_c, up_s), + gate=(gt_c, gt_s), + down=(dn_c, dn_s), + ) + return _STACK + + +def _router(gen: torch.Generator): + """Router weight and bias. The weight sits on a 2^-8 grid within +-1/32 so + that, against a router input on a quarter grid, every partial sum of a + router GEMV is exact in fp32: any summation order -- the kernel's, torch's + -- gives the same logits.""" + weight = ( + torch.randint(-8, 9, (EXPERTS, HIDDEN), device=DEV, generator=gen).float() * 2.0**-8 + ).to(torch.bfloat16) + bias = (torch.randn(EXPERTS, device=DEV, generator=gen) * 0.05).to(torch.bfloat16) + return weight, bias + + +def _router_input(tokens: int, gen: torch.Generator) -> torch.Tensor: + """bf16 router input on a quarter grid in [-1, 1] (see `_router`).""" + ints = torch.randint(-4, 5, (tokens, HIDDEN), device=DEV, generator=gen) + return (ints.float() * 0.25).to(torch.bfloat16) + + +def _input_global_scale(x: torch.Tensor) -> torch.Tensor: + """The FC1-input global scale a checkpoint's `input_scale` encodes.""" + return (E4M3_MAX * E2M1_MAX / x.float().abs().amax()).reshape(1).to(torch.float32) + + +def _quantize(x: torch.Tensor, g1: torch.Tensor): + """The experts' NVFP4 input, as the target quantizes it: linear scales.""" + xq, xsf = torch.ops.trtllm.fp4_quantize(x, g1, SV, False, False) + return xq, xsf.view(torch.float8_e4m3fn) + + +def _scalars(xq, xsf, chosen, g1: torch.Tensor): + """The three per-expert scalars for weights with a global scale of 1, and + `g2` calibrated on the true FC1 activation of the chosen experts.""" + stack = _stack() + x = _dequant_act(xq, xsf) / g1 + amax = torch.zeros((), device=DEV) + with true_fp32_matmul(): + for e in chosen.unique().tolist(): + tok = (chosen == e).any(dim=1).nonzero(as_tuple=True)[0] + up = x[tok] @ _dequant(*_slice(stack["up"], e)).t() + gate = x[tok] @ _dequant(*_slice(stack["gate"], e)).t() + amax = torch.maximum(amax, (up * gate * torch.sigmoid(gate)).abs().max()) + g2 = E4M3_MAX * E2M1_MAX / amax + a1 = torch.full((EXPERTS,), 1.0, device=DEV) / g1 + return a1 * g2, a1, torch.full((EXPERTS,), 1.0, device=DEV) / g2 + + +_E2M1 = None + + +def _dequant(codes: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + global _E2M1 + if _E2M1 is None: + _E2M1 = torch.tensor( + [0, 0.5, 1, 1.5, 2, 3, 4, 6, -0.0, -0.5, -1, -1.5, -2, -3, -4, -6], + dtype=torch.float32, + device=DEV, + ) + return _E2M1[codes.long()] * scales.float().repeat_interleave(SV, dim=1) + + +def _slice(operand, e: int): + codes, scales = operand + return codes[e], scales[e] + + +def _dequant_act(xq: torch.Tensor, xsf: torch.Tensor) -> torch.Tensor: + codes = torch.stack([xq & 0x0F, xq >> 4], dim=-1).reshape(xq.shape[0], -1) + return _dequant(codes, xsf.view(torch.float8_e4m3fn).reshape(xq.shape[0], -1)) + + +def _selection(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor): + """`(logits, top-8 by score + bias, smallest 8th-to-9th score gap)`.""" + with true_fp32_matmul(): + logits = x.float() @ weight.float().t() + sel = 0.5 * torch.tanh(0.5 * logits) + 0.5 + bias.float() + top = torch.topk(sel, TOP_K + 1, dim=-1) + margin = (top.values[:, TOP_K - 1] - top.values[:, TOP_K]).min().item() + return logits, top.indices[:, :TOP_K], margin + + +def _case(tokens: int, seed: int, clear_margin: bool = False) -> dict: + """One call's inputs. `clear_margin` redraws (deterministically, from the + next seeds) a router input whose expert choice sits within + `SELECTION_MARGIN` of a tie, which only a comparison against the + reference needs.""" + for attempt in range(16): + gen = torch.Generator(device=DEV).manual_seed(seed + attempt) + weight, bias = _router(gen) + x = _router_input(tokens, gen) + logits, chosen, margin = _selection(x, weight, bias) + if not clear_margin or margin > SELECTION_MARGIN: + break + else: + raise AssertionError(f"no draw from seed {seed} clears the selection margin") + g1 = _input_global_scale(x) + xq, xsf = _quantize(x, g1) + s1, sg, s2 = _scalars(xq, xsf, chosen, g1) + return dict( + x=x, + g1=g1, + xq=xq, + xsf=xsf, + logits=logits, + router_weight=weight, + routing_bias=bias, + output1_scale_scalar=s1, + output1_scale_gate_scalar=sg, + output2_scale_scalar=s2, + ) + + +# ── the call ────────────────────────────────────────────────────────────── + + +def _call( + case: dict, + *, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + fused_router: bool, + workspace=None, + event=None, + gemm2_k128: bool, + do_finalize: bool = False, + routing_logits=None, +): + """The target's call, every argument explicit.""" + stack = _stack() + tokens = case["x"].shape[0] + if routing_logits is None: + # With the fused router only the shape is read; filled so that an + # autotuning run would profile on plausible routing. + routing_logits = ( + torch.randn(tokens, EXPERTS, device=DEV) if fused_router else case["logits"] + ) + return op( + routing_logits=routing_logits, + routing_bias=case["routing_bias"], + hidden_states=hidden_states, + hidden_states_scale=hidden_states_scale, + gemm1_weights=stack["gemm1_weights"], + gemm1_weights_scale=stack["gemm1_weights_scale"], + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, + gemm2_weights=stack["gemm2_weights"], + gemm2_weights_scale=stack["gemm2_weights_scale"], + gemm2_bias=None, + output1_scale_scalar=case["output1_scale_scalar"], + output1_scale_gate_scalar=case["output1_scale_gate_scalar"], + output2_scale_scalar=case["output2_scale_scalar"], + num_experts=EXPERTS, + top_k=TOP_K, + n_group=1, + topk_group=1, + intermediate_size=INTER, + local_expert_offset=0, + local_num_experts=EXPERTS, + routed_scaling_factor=ROUTED_SCALE, + routing_method_type=ROUTING_METHOD_DEEPSEEK_V3, + do_finalize=do_finalize, + act_type=ACT_TYPE_SWIGLU, + topk_weights=None, + topk_ids=None, + output=None, + tune_max_num_tokens=8192, + use_dp=False, + hidden_states_ready_event=None if event is None else event.cuda_event, + router_hidden_states=case["x"] if fused_router else None, + router_weight=case["router_weight"] if fused_router else None, + router_workspace=workspace if fused_router else None, + prefer_gemm2_k128=gemm2_k128, + ) + + +def _quantize_aside(case: dict, side: torch.cuda.Stream, delay: bool, out=None): + """The target's input path: quantize on `side`, forked after the current + stream's work, and return `(xq, xsf, event)` with `event` recorded there. + `out` holds buffers to write instead of fresh tensors; `delay` sleeps on + `side` first.""" + forked = torch.cuda.Event() + ready = torch.cuda.Event() + forked.record() + with torch.cuda.stream(side): + side.wait_event(forked) + if delay: + torch.cuda._sleep(SLEEP_CYCLES) + xq, xsf = _quantize(case["x"], case["g1"]) + if out is not None: + out[0].copy_(xq) + out[1].copy_(xsf) + xq, xsf = out + ready.record(side) + return xq, xsf, ready + + +def _valid_rows(outs): + """The unfinalized outputs as `(rows each slot reads, weights, map)`.""" + rows, weights, perm = outs + return rows[perm.long().reshape(-1)], weights, perm + + +def _finalize(outs) -> torch.Tensor: + """`sum_k weights[t, k] * rows[map[t, k]]` in fp32, rounded once: what the + post-MoE hop computes.""" + rows, weights, perm = outs + gathered = rows.float()[perm.long()] + return (gathered * weights.float().unsqueeze(-1)).sum(dim=1).to(torch.bfloat16) + + +def _assert_same(a, b, what: str) -> None: + for x, y, name in zip(_valid_rows(a), _valid_rows(b), ("rows", "weights", "map")): + assert torch.equal(x, y), f"{what}: {name} differ" + + +def _kernels(fn) -> list[str]: + """Names of the kernels `fn` launches, in the order they start.""" + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + fn() + torch.cuda.synchronize() + launched = [ + e + for e in prof.events() + if str(e.device_type) == "DeviceType.CUDA" and e.device_time_total > 0 + ] + return [e.name for e in sorted(launched, key=lambda e: e.time_range.start)] + + +def _k128_covers(tokens: int) -> bool: + """gemm2K128's gate at tile N 8: one CTA per SM, at most 15 units (token + tile x 128 output features) each, at most 64 token tiles.""" + tiles = min(tokens * TOP_K, EXPERTS) + sms = torch.cuda.get_device_properties(0).multi_processor_count + units = tiles * HIDDEN // 128 + return tiles <= 64 and -(-units // sms) <= 15 + + +# ── tests ───────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 + """Each cell as the target makes the call, against the reference. + + The late-join cells quantize on a side stream behind a sleep, so the + result also proves FC1 waited for the input. + """ + certifying(op) + spec = cell.spec + geometry = (spec["hidden"], spec["intermediate"], spec["num_experts"], spec["top_k"]) + assert geometry == (HIDDEN, INTER, EXPERTS, TOP_K), geometry + assert spec["routed_scaling_factor"] == ROUTED_SCALE and spec["local_expert_offset"] == 0 + case = _case(spec["tokens"], seed=zlib.crc32(cell.why.encode()) % 2**31, clear_margin=True) + workspace = make_fused_router_workspace(DEV) if spec["fused_router"] else None + if spec["late_join"]: + side = torch.cuda.Stream() + xq, xsf, ready = _quantize_aside(case, side, delay=True) + else: + xq, xsf, ready = case["xq"], case["xsf"], None + outs = _call( + case, + hidden_states=xq, + hidden_states_scale=xsf, + fused_router=spec["fused_router"], + workspace=workspace, + event=ready, + gemm2_k128=spec["gemm2_k128"], + ) + torch.cuda.synchronize() + + stack = _stack() + ref, _, ref_weights = op.reference( + routing_logits=None if spec["fused_router"] else case["logits"], + routing_bias=case["routing_bias"], + hidden_states=case["xq"], + hidden_states_scale=case["xsf"], + up=stack["up"], + gate=stack["gate"], + down=stack["down"], + output1_scale_scalar=case["output1_scale_scalar"], + output1_scale_gate_scalar=case["output1_scale_gate_scalar"], + output2_scale_scalar=case["output2_scale_scalar"], + top_k=TOP_K, + local_expert_offset=0, + routed_scaling_factor=ROUTED_SCALE, + router_hidden_states=case["x"] if spec["fused_router"] else None, + router_weight=case["router_weight"] if spec["fused_router"] else None, + ) + # Each token's weights, as a set: the kernel stores them in its top-k order. + torch.testing.assert_close( + outs[1].float().sort(dim=-1).values, + ref_weights.float().sort(dim=-1).values, + rtol=2.0**-7, + atol=0.0, + ) + _compare(_finalize(outs), ref) + + +@pytest.mark.parametrize("tokens", range(1, 9)) +def test_fused_router_matches_the_router_gemm(tokens) -> None: + """The fused router computes the logits of `dsv3_router_gemm_op` bit for + bit and routes them as the routing kernel does: every output is bitwise + the one the op gives when routed from those logits. Over every row count + the kernel has an instance for, on an off-grid router input where + summation order shows.""" + gen = torch.Generator(device=DEV).manual_seed(4000 + tokens) + case = _case(tokens, seed=4100 + tokens) + x = torch.randn(tokens, HIDDEN, device=DEV, generator=gen).to(torch.bfloat16) + weight = (torch.randn(EXPERTS, HIDDEN, device=DEV, generator=gen) * 0.02).to(torch.bfloat16) + case.update(x=x, router_weight=weight) + logits = torch.ops.trtllm.dsv3_router_gemm_op(x, weight.t(), None, torch.float32) + workspace = make_fused_router_workspace(DEV) + common = dict(hidden_states=case["xq"], hidden_states_scale=case["xsf"], gemm2_k128=False) + fused = _call(case, fused_router=True, workspace=workspace, **common) + routed = _call(case, fused_router=False, routing_logits=logits, **common) + _assert_same(fused, routed, f"{tokens} rows") + + +@pytest.mark.parametrize("tokens", (1, 8)) +def test_gemm2_k128_matches_the_cubin(tokens) -> None: + """FC2 at K = 128 writes the Gemm2 cubin's rows, bitwise, and runs where + its gate says it does.""" + case = _case(tokens, seed=4200 + tokens) + common = dict(hidden_states=case["xq"], hidden_states_scale=case["xsf"], fused_router=False) + k128 = _call(case, gemm2_k128=True, **common) + cubin = _call(case, gemm2_k128=False, **common) + _assert_same(k128, cubin, f"{tokens} rows") + + names = _kernels(lambda: _call(case, gemm2_k128=True, **common)) + ran = any("gemm2K128Kernel" in n for n in names) + assert ran == _k128_covers(tokens), (names, _k128_covers(tokens)) + names = _kernels(lambda: _call(case, gemm2_k128=False, **common)) + assert not any("gemm2K128Kernel" in n for n in names), names + + +def test_c1_call_launches_three_kernels() -> None: + """The c1 call is the fused router, FC1 and FC2: no router GEMV, no + separate routing kernel.""" + case = _case(8, seed=4300) + workspace = make_fused_router_workspace(DEV) + shape_only = torch.randn(8, EXPERTS, device=DEV) + names = _kernels( + lambda: _call( + case, + hidden_states=case["xq"], + hidden_states_scale=case["xsf"], + fused_router=True, + workspace=workspace, + gemm2_k128=True, + routing_logits=shape_only, + ) + ) + assert len(names) == 3, names + assert "routingFusedRouterKernel" in names[0], names + assert ("gemm2K128Kernel" in names[2]) == _k128_covers(8), names + + +def test_late_join_waits_for_the_input() -> None: + """FC1 reads the input only after the event: the buffers handed to the op + hold a different activation's quantization until the side stream, held + back by a sleep, overwrites them. Routing launches without waiting.""" + case = _case(8, seed=4400) + stale = _case(8, seed=4401) + workspace = make_fused_router_workspace(DEV) + expected = _call( + case, + hidden_states=case["xq"], + hidden_states_scale=case["xsf"], + fused_router=True, + workspace=workspace, + gemm2_k128=True, + ) + buffers = (stale["xq"].clone(), stale["xsf"].clone()) + side = torch.cuda.Stream() + xq, xsf, ready = _quantize_aside(case, side, delay=True, out=buffers) + late = _call( + case, + hidden_states=xq, + hidden_states_scale=xsf, + fused_router=True, + workspace=workspace, + event=ready, + gemm2_k128=True, + ) + torch.cuda.synchronize() + _assert_same(late, expected, "late join") + + +def test_c1_call_replays_in_a_cuda_graph() -> None: + """The c1 step captured as the target runs it -- input quantized on a side + stream, joined inside the op -- replays the eager result on new inputs, and + every launch leaves the router workspace ready for the next one.""" + cases = [_case(8, seed=4500 + i) for i in range(3)] + static = {k: v.clone() for k, v in cases[0].items()} + workspace = make_fused_router_workspace(DEV) + side = torch.cuda.Stream() + + def step(): + xq, xsf, ready = _quantize_aside(static, side, delay=False) + return _call( + static, + hidden_states=xq, + hidden_states_scale=xsf, + fused_router=True, + workspace=workspace, + event=ready, + gemm2_k128=True, + routing_logits=static["logits"], + ) + + step() # load the kernels before capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = step() + + launches = 1 + for case in cases: + for key, value in case.items(): + static[key].copy_(value) + graph.replay() + launches += 1 + torch.cuda.synchronize() + eager = _call( + case, + hidden_states=case["xq"], + hidden_states_scale=case["xsf"], + fused_router=True, + workspace=workspace, + gemm2_k128=True, + ) + launches += 1 + torch.cuda.synchronize() + _assert_same(captured, eager, "replay") + words = workspace.view(torch.int32) + epoch = int(words[_EPOCH_WORD].item()) + assert epoch == launches, (epoch, launches) + current = words[(epoch & 1) * _BUFFER_WORDS : ((epoch & 1) + 1) * _BUFFER_WORDS] + assert bool((current == -(2**31)).all()), "the next launch's exchange buffer is not idle" + + +def test_fused_router_refuses_more_than_8_rows() -> None: + """Past 8 rows the fused router is a loud error; the caller routes from + the router GEMM's logits instead (the 32-row cell).""" + case = _case(9, seed=4600) + with pytest.raises(RuntimeError, match="fused router"): + _call( + case, + hidden_states=case["xq"], + hidden_states_scale=case["xsf"], + fused_router=True, + workspace=make_fused_router_workspace(DEV), + gemm2_k128=True, + ) + + +# ── guards ──────────────────────────────────────────────────────────────── + + +def test_guard_refuses_a_strided_router_input() -> None: + case = _case(2, seed=4700) + case["x"] = torch.cat([case["x"], case["x"]], dim=1)[:, ::2] + with pytest.raises(AssertionError, match="router_hidden_states must be contiguous"): + with validating(op): + _call( + case, + hidden_states=case["xq"], + hidden_states_scale=case["xsf"], + fused_router=True, + workspace=make_fused_router_workspace(DEV), + gemm2_k128=True, + ) + + +def test_guard_refuses_a_zero_filled_workspace() -> None: + """A zero-filled workspace reads as published zero sums: routing on zero + logits, with no error.""" + case = _case(2, seed=4701) + workspace = torch.zeros(FUSED_ROUTER_WORKSPACE_WORDS, dtype=torch.int32, device=DEV) + with pytest.raises(AssertionError, match="between-launch state"): + with validating(op): + _call( + case, + hidden_states=case["xq"], + hidden_states_scale=case["xsf"], + fused_router=True, + workspace=workspace, + gemm2_k128=True, + ) + + +def test_guard_refuses_an_unrecorded_event() -> None: + """An event never recorded has no cudaEvent_t yet: FC1 would not wait.""" + case = _case(2, seed=4702) + with pytest.raises(AssertionError, match="null cudaEvent_t"): + with validating(op): + _call( + case, + hidden_states=case["xq"], + hidden_states_scale=case["xsf"], + fused_router=True, + workspace=make_fused_router_workspace(DEV), + event=torch.cuda.Event(), + gemm2_k128=True, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_op_matrix.py new file mode 100644 index 000000000000..096da8663092 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_op_matrix.py @@ -0,0 +1,622 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for the two collective MTP draft-step entries of the GLM-5.2 tp16 target. + +`glm_draft_ehproj_table` and `glm_draft_sample_glue` exchange through an MNNVL workspace across the TP group, +so neither can run in one process. Run plainly, this script is its own launcher: it spawns one rank per device +named in CUDA_VISIBLE_DEVICES, under a deadline it enforces by killing every rank's session, because a broken +exchange hangs (a rank polls for a packet that never lands) rather than raising. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _draft_step_op_matrix.py draft_ehproj_table + +Each rank drives every cell of the entry, a draft loop of the target's shape (six table steps and the last glue +step, repeated) eagerly and as CUDA-graph replays, and the entry's guards. The 16-rank geometry of the table +kernel (a CTA's columns spanning two ranks' hidden-half shares) needs 16 ranks over 4 nodes: start one process +per GPU with an external launcher and `--rank-worker`, with RANK, WORLD_SIZE, LOCAL_RANK, MASTER_ADDR and +MASTER_PORT in the environment (SLURM_PROCID, SLURM_NTASKS and SLURM_LOCALID stand in for the first three): + + srun -N4 --ntasks-per-node=4 python _draft_step_op_matrix.py draft_ehproj_table --rank-worker + +Not a pytest module: the collected entry points are `test_modeling_v2_draft_ehproj_table.py` and +`test_modeling_v2_draft_sample_glue.py` beside this file. The launcher half imports nothing from +tensorrt_llm, and the ranks reach the target's ops by absolute import. +""" + +import os +import signal +import socket +import subprocess +import sys +import time +import zlib +from pathlib import Path +from typing import Any, Dict, List + +import torch +import torch.distributed as dist + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 1500 + +HIDDEN = 6144 +EPS = 1e-5 +MAX_BLOCKS = 64 +#: The DSA indexer-K cache layout of GLM-5.2: 64 tokens per block, 128 FP8 bytes and one FP32 scale per token. +TOKENS_PER_BLOCK = 64 +DATA_BYTES = 128 +SCALE_BYTES = 4 +TOKEN_LANES = 8 + +# Bound in the rank body: importing the entries pulls in tensorrt_llm. +RANK = 0 +WORLD = 1 +HOST_GROUP: Any = None +table_op: Any = None +glue_op: Any = None +validating: Any = None +assert_within_ulp: Any = None +WORKSPACES: Dict[str, Any] = {} + + +def _generator(*key: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(zlib.crc32(repr(key).encode()) % 2**31) + + +def _logits(spec: dict, seed: int, rank: int) -> torch.Tensor: + """Rank `rank`'s BF16 logits shard; every rank can rebuild every other rank's.""" + rows, shard = spec["rows"], spec["vocab_shard"] + g = _generator(seed, 1, rank) + logits = torch.randn(rows, shard, generator=g, device="cuda") + if spec["logits"] == "tied": + # Every rank's maximum is 8 at two local positions: the lower position of rank 0 wins. + first = torch.randint(0, shard // 2, (rows,), generator=g, device="cuda") + second = first + shard // 2 + rows_index = torch.arange(rows, device="cuda") + logits[rows_index, first] = 8.0 + logits[rows_index, second] = 8.0 + return logits.to(torch.bfloat16) + + +class _LazyTables: + """Every rank's token table, built on first use: the reference reads only the winners' rows.""" + + def __init__(self, seed: int, shard: int): + self._seed, self._shard, self._tables = seed, shard, {} + + def __getitem__(self, rank: int) -> torch.Tensor: + if rank not in self._tables: + g = _generator(self._seed, 2, rank) + self._tables[rank] = torch.randn( + self._shard, HIDDEN, generator=g, device="cuda", dtype=torch.bfloat16 + ) + return self._tables[rank] + + def __len__(self) -> int: + return WORLD + + +def _hidden_half(seed: int, rank: int, rows: int) -> torch.Tensor: + g = _generator(seed, 3, rank) + return torch.randn(rows, HIDDEN // WORLD, generator=g, device="cuda") + + +def _state(spec: dict, seed: int) -> Dict[str, Any]: + """The step's state arguments, identical on every rank.""" + rows, draft_len, step = spec["rows"], spec["draft_len"], spec["step"] + g = _generator(seed, 4) + vocab = spec["vocab_shard"] * WORLD + num_accepted = torch.randint(1, draft_len + 2, (rows,), generator=g, device="cuda").int() + batch_indices = torch.randperm(rows, generator=g, device="cuda").int() + state: Dict[str, Any] = dict( + draft_tokens=torch.randint(0, vocab, (rows, draft_len), generator=g, device="cuda").int(), + draft_step=step, + positions_in=None, + positions_out=None, + gather_ids=None, + seq_lens=None, + fill_seq_lens=spec["fill"], + kv_lens=None, + kv_update=spec["kv_update"], + rewind_accepted=None, + draft_len=draft_len, + block_offsets=None, + dsa_layout=[], + req_idx_per_token=None, + slot_mapping_fp8=None, + slot_mapping_scale=None, + gen_kv_indptr=None, + gen_cached_token_indptr=None, + ) + verify_rows = rows * (draft_len + 1) + if spec["positions"] == "gather": + state["positions_in"] = torch.randint( + 100, 100000, (verify_rows,), generator=g, device="cuda" + ).int() + state["gather_ids"] = batch_indices.long() * (draft_len + 1) + num_accepted.long() - 1 + state["positions_out"] = torch.full((rows,), -9, dtype=torch.int32, device="cuda") + elif spec["positions"] == "plain": + state["positions_in"] = torch.randint( + 100, 100000, (rows,), generator=g, device="cuda" + ).int() + state["positions_out"] = torch.full((rows,), -9, dtype=torch.int32, device="cuda") + seq_before = draft_len + 1 if spec["fill"] else 1 + state["seq_lens"] = torch.full((rows,), seq_before, dtype=torch.int32, device="cuda") + state["kv_lens"] = ( + 1000 + 37 * torch.arange(rows, device="cuda", dtype=torch.int32) + 3 * step + ).int() + if spec["kv_update"] == 2: + state["rewind_accepted"] = num_accepted + if spec["dsa"]: + num_tokens = rows + state["block_offsets"] = torch.randint( + 0, 4096, (rows, MAX_BLOCKS), generator=g, device="cuda" + ).int() + state["dsa_layout"] = [ + num_tokens, + TOKEN_LANES, + TOKENS_PER_BLOCK, + DATA_BYTES, + SCALE_BYTES, + TOKENS_PER_BLOCK * (DATA_BYTES + SCALE_BYTES), + TOKENS_PER_BLOCK * DATA_BYTES, + ] + state["req_idx_per_token"] = torch.full((num_tokens + 3,), -5, dtype=torch.int32).cuda() + state["slot_mapping_fp8"] = torch.full((num_tokens + 3,), -5, dtype=torch.int64).cuda() + state["slot_mapping_scale"] = torch.full((num_tokens + 3,), -5, dtype=torch.int64).cuda() + state["gen_kv_indptr"] = torch.full((rows + 4,), -5, dtype=torch.int64).cuda() + state["gen_cached_token_indptr"] = torch.full((rows + 4,), -5, dtype=torch.int64).cuda() + state["_num_accepted"] = num_accepted + state["_batch_indices"] = batch_indices + return state + + +def _snapshot(state: Dict[str, Any]) -> Dict[str, Any]: + return { + k: (v.clone() if isinstance(v, torch.Tensor) else v) + for k, v in state.items() + if not k.startswith("_") + } + + +def _assert_exact(name: str, out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + try: + torch.testing.assert_close(out, ref, rtol=0, atol=0) + except AssertionError as exc: + raise AssertionError(f"{where}: {name}: {exc}") from exc + + +def _compare( + expected: Dict[str, torch.Tensor], actual: Dict[str, torch.Tensor], where: str +) -> None: + """Integers and x bit for bit; the normalized rows within one RMSNorm's rounding. + + The RMSNorm outputs (`normed_out`, `glue_out`) differ from the FP32 reference in the order of the 6144 + squares and in the kernel's approximate reciprocal square root: a few 1e-7 relative on the scale, one BF16 + ulp on the rare element that close to a rounding midpoint, which is up to twice the band's unit (2**-8 of the + row maximum) in the maximum's binade. So 2 ulp against the row scale and 0.25 ulp of relative RMS, as for the + step-0 glue kernel. + """ + for name, ref in expected.items(): + out = actual[name] + if name in ("normed_out", "glue_out"): + try: + assert_within_ulp(out, ref, element_ulp=2.0, rms_ulp=0.25) + except AssertionError as exc: + raise AssertionError(f"{where}: {name}: {exc}") from exc + else: + _assert_exact(name, out, ref, where) + + +def _table_call(spec: dict, seed: int, logits: List[torch.Tensor], tables: _LazyTables): + """One table-op call on this rank: (call kwargs, reference kwargs).""" + rows = spec["rows"] + state = _state(spec, seed) + halves = [_hidden_half(seed, r, rows) for r in range(WORLD)] + norm_g = _generator(seed, 5) + norm_weight = (1.0 + 0.1 * torch.randn(HIDDEN, generator=norm_g, device="cuda")).to( + torch.bfloat16 + ) + workspace = WORKSPACES["table"] + kwargs = dict( + logits=logits[RANK], + exchange_buffer=workspace.buffer, + exchange_flags=workspace.flags, + group_size=WORLD, + hidden_half=halves[RANK], + token_table=tables[RANK], + norm_weight=norm_weight, + eps=EPS, + normed_out=torch.full((rows, HIDDEN), 7.0, dtype=torch.bfloat16, device="cuda"), + residual_out=torch.full((rows, HIDDEN), 7.0, dtype=torch.bfloat16, device="cuda"), + **{k: v for k, v in state.items() if not k.startswith("_")}, + ) + ref_kwargs = dict( + kwargs, **_snapshot(state), logits=logits, hidden_half=halves, token_table=tables + ) + return kwargs, ref_kwargs + + +def _glue_call(spec: dict, seed: int, logits: List[torch.Tensor], embed_table: torch.Tensor): + """One glue-op call on this rank: (call kwargs, reference kwargs).""" + rows, draft_len = spec["rows"], spec["draft_len"] + state = _state(spec, seed) + g = _generator(seed, 6) + window = 2 * HIDDEN // WORLD + kwargs = dict( + logits=logits[RANK], + exchange_buffer=WORKSPACES["glue"].buffer, + exchange_flags=WORKSPACES["glue"].flags, + group_size=WORLD, + embed_table=None, + hidden=None, + enorm_weight=None, + hnorm_weight=None, + eps=EPS, + col_lo=0, + col_hi=0, + glue_out=None, + accepted_tokens=None, + batch_indices=None, + num_accepted=None, + next_new_tokens=None, + **{k: v for k, v in state.items() if not k.startswith("_")}, + ) + if spec["glue"]: + kwargs.update( + embed_table=embed_table, + hidden=(3.0 * torch.randn(rows, HIDDEN, generator=g, device="cuda")).to(torch.bfloat16), + enorm_weight=(1.0 + 0.1 * torch.randn(HIDDEN, generator=g, device="cuda")).to( + torch.bfloat16 + ), + hnorm_weight=(1.0 + 0.1 * torch.randn(HIDDEN, generator=g, device="cuda")).to( + torch.bfloat16 + ), + col_lo=RANK * window, + col_hi=(RANK + 1) * window, + glue_out=torch.full((rows, window), 7.0, dtype=torch.bfloat16, device="cuda"), + ) + if spec["next_new"]: + kwargs.update( + accepted_tokens=torch.randint( + 0, 1000, (rows, draft_len + 1), generator=g, device="cuda" + ).int(), + batch_indices=state["_batch_indices"], + num_accepted=state["_num_accepted"], + next_new_tokens=torch.full((rows, draft_len + 1), -3, dtype=torch.int32, device="cuda"), + ) + ref_kwargs = dict(kwargs, **_snapshot(state), logits=logits) + return kwargs, ref_kwargs + + +_TABLE_OUTPUTS = ("draft_tokens", "normed_out", "residual_out") +_GLUE_OUTPUTS = ("draft_tokens", "glue_out", "next_new_tokens") +_STATE_OUTPUTS = ( + "positions_out", + "seq_lens", + "kv_lens", + "req_idx_per_token", + "slot_mapping_fp8", + "slot_mapping_scale", + "gen_kv_indptr", + "gen_cached_token_indptr", +) + + +def _written(kwargs: Dict[str, Any], names) -> Dict[str, torch.Tensor]: + return {n: kwargs[n] for n in names if kwargs.get(n) is not None} + + +def check_table_cells() -> int: + for index, cell in enumerate(table_op.CELLS): + spec = cell.spec + seed = zlib.crc32(cell.why.encode()) % 2**31 + logits = [_logits(spec, seed, r) for r in range(WORLD)] + tables = _LazyTables(seed, spec["vocab_shard"]) + kwargs, ref_kwargs = _table_call(spec, seed, logits, tables) + with validating(table_op): + table_op(**kwargs) + torch.cuda.synchronize() + expected = table_op.reference(**ref_kwargs) + _compare(expected, _written(kwargs, _TABLE_OUTPUTS + _STATE_OUTPUTS), f"cell {index}") + return len(table_op.CELLS) + + +def check_glue_cells() -> int: + embed_table = None + for index, cell in enumerate(glue_op.CELLS): + spec = cell.spec + seed = zlib.crc32(cell.why.encode()) % 2**31 + if spec["glue"] and embed_table is None: + embed_table = torch.randn( + spec["vocab_shard"] * WORLD, + HIDDEN, + generator=_generator(9), + device="cuda", + dtype=torch.bfloat16, + ) + logits = [_logits(spec, seed, r) for r in range(WORLD)] + kwargs, ref_kwargs = _glue_call(spec, seed, logits, embed_table) + with validating(glue_op): + glue_op(**kwargs) + torch.cuda.synchronize() + expected = glue_op.reference(**ref_kwargs) + _compare(expected, _written(kwargs, _GLUE_OUTPUTS + _STATE_OUTPUTS), f"cell {index}") + return len(glue_op.CELLS) + + +def _draft_loop_specs(draft_len: int = 7): + """The target's c1 draft loop: table steps 0..K-2, then the glue for step K-1.""" + table_steps = [] + for step in range(draft_len - 1): + table_steps.append( + table_op.CELLS[0].spec if step == 0 else dict(table_op.CELLS[1].spec, step=step) + ) + glue_step = dict(glue_op.CELLS[0].spec, step=draft_len - 1) + return table_steps, glue_step + + +def check_draft_loop_eager_and_graph() -> None: + """Seven consecutive calls per decode step over two workspaces, eagerly and as graph replays. + + The Lamport flags rotate on every call and each call empties the buffer of the previous one, so the + exchange must stay correct across steps, decode steps and replays of a captured loop. + """ + table_steps, glue_step = _draft_loop_specs() + calls = [] + for step, spec in enumerate(table_steps): + seed = 1000 + step + logits = [_logits(spec, seed, r) for r in range(WORLD)] + tables = _LazyTables(seed, spec["vocab_shard"]) + calls.append(("table", _table_call(spec, seed, logits, tables))) + seed = 1000 + len(table_steps) + logits = [_logits(glue_step, seed, r) for r in range(WORLD)] + calls.append(("glue", _glue_call(glue_step, seed, logits, None))) + + def run_loop() -> None: + for kind, (kwargs, _) in calls: + (table_op if kind == "table" else glue_op)(**kwargs) + + def reset_and_expect(): + """Restore every buffer the loop updates in place; return the expected outputs per call.""" + expected = [] + for kind, (kwargs, ref_kwargs) in calls: + for name, value in ref_kwargs.items(): + # The snapshots of the buffers the op updates in place are separate tensors. + live = kwargs.get(name) + if isinstance(value, torch.Tensor) and isinstance(live, torch.Tensor): + if live is not value: + live.copy_(value) + op = table_op if kind == "table" else glue_op + expected.append(op.reference(**ref_kwargs)) + return expected + + for repeat in range(3): + expected = reset_and_expect() + run_loop() + torch.cuda.synchronize() + for (kind, (kwargs, _)), exp in zip(calls, expected): + names = (_TABLE_OUTPUTS if kind == "table" else _GLUE_OUTPUTS) + _STATE_OUTPUTS + _compare(exp, _written(kwargs, names), f"eager loop {repeat} {kind}") + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(side): + with torch.cuda.graph(graph, stream=side): + run_loop() + torch.cuda.current_stream().wait_stream(side) + for repeat in range(3): + expected = reset_and_expect() + graph.replay() + torch.cuda.synchronize() + for (kind, (kwargs, _)), exp in zip(calls, expected): + names = (_TABLE_OUTPUTS if kind == "table" else _GLUE_OUTPUTS) + _STATE_OUTPUTS + _compare(exp, _written(kwargs, names), f"graph replay {repeat} {kind}") + + +def check_guards_refuse() -> None: + """The guards refuse their silent-wrong inputs on every rank before anything is launched.""" + spec = dict(table_op.CELLS[0].spec) + seed = 77 + logits = [_logits(spec, seed, r) for r in range(WORLD)] + kwargs, _ = _table_call(spec, seed, logits, _LazyTables(seed, spec["vocab_shard"])) + bad = dict(kwargs, positions_out=kwargs["positions_in"]) + try: + with validating(table_op): + table_op(**bad) + except AssertionError as exc: + assert "must not overlap" in str(exc), exc + else: + raise AssertionError("the table guard accepted positions_out aliasing positions_in") + bad = dict(kwargs, draft_len=kwargs["draft_len"] + 1) + try: + with validating(table_op): + table_op(**bad) + except AssertionError as exc: + assert "KV rewind" in str(exc), exc + else: + raise AssertionError("the table guard accepted a rewind with another draft length") + spec = dict(glue_op.CELLS[0].spec) + logits = [_logits(spec, seed, r) for r in range(WORLD)] + kwargs, _ = _glue_call(spec, seed, logits, None) + wide = torch.zeros(spec["rows"], spec["draft_len"] + 2, dtype=torch.int32, device="cuda") + try: + with validating(glue_op): + glue_op(**dict(kwargs, next_new_tokens=wide)) + except AssertionError as exc: + assert "next_new_tokens must be exactly" in str(exc), exc + else: + raise AssertionError("the glue guard accepted a wider next_new_tokens") + + +def _rank_env(): + env = os.environ + rank = int(env.get("RANK", env.get("SLURM_PROCID", "0"))) + world = int(env.get("WORLD_SIZE", env.get("SLURM_NTASKS", "1"))) + local_rank = int(env.get("LOCAL_RANK", env.get("SLURM_LOCALID", str(rank)))) + return rank, world, local_rank + + +def _run_one_rank(entry: str, expected_cells: int) -> int: + """Body of one rank: build the workspaces, run the entry's checks, exit hard on a failure.""" + global RANK, WORLD, HOST_GROUP, table_op, glue_op, validating, assert_within_ulp + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank, world, local_rank = _rank_env() + os.environ.setdefault("RANK", str(rank)) + os.environ.setdefault("WORLD_SIZE", str(world)) + torch.cuda.set_device(local_rank) + dist.init_process_group( + backend="cuda:nccl,cpu:gloo", + init_method="env://", + rank=rank, + world_size=world, + device_id=torch.device("cuda", local_rank), + ) + RANK, WORLD = rank, world + HOST_GROUP = dist.new_group(backend="gloo") + + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec import ( # noqa: E501 + _draft_exchange, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec.draft_ehproj_table import ( # noqa: E501 + glm_draft_ehproj_table, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec.draft_sample_glue import ( # noqa: E501 + glm_draft_sample_glue, + ) + + table_op, glue_op = glm_draft_ehproj_table, glm_draft_sample_glue + # Spawned, not collected: reach the tree's _validating the plain way. tests/unittest/_torch/modeling_v2 is + # four levels above this directory. + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import assert_within_ulp as _within + from _validating import validating as _validating_cm + + validating, assert_within_ulp = _validating_cm, _within + + hosts: List[Any] = [None] * WORLD + dist.all_gather_object(hosts, socket.gethostname(), group=HOST_GROUP) + fabric = len(set(hosts)) > 1 or os.environ.get("TRTLLM_FORCE_MNNVL_AR", "0") == "1" + assert _draft_exchange.ehproj_table_supports_tp(WORLD), ( + f"the table op does not serve {WORLD} ranks" + ) + max_rows = _draft_exchange.MAX_ROWS + WORKSPACES["table"] = _draft_exchange.DraftExchangeWorkspace( + max_rows * WORLD * _draft_exchange.ehproj_table_slot_bytes(WORLD), + WORLD, + RANK, + local_rank, + fabric, + dist.group.WORLD, + ) + WORKSPACES["glue"] = _draft_exchange.DraftExchangeWorkspace( + max_rows * WORLD * _draft_exchange.sample_glue_slot_bytes(WORLD), + WORLD, + RANK, + local_rank, + fabric, + dist.group.WORLD, + ) + + checks = { + "draft_ehproj_table": (check_table_cells,), + "draft_sample_glue": (check_glue_cells,), + }[entry] + (check_draft_loop_eager_and_graph, check_guards_refuse) + for check in checks: + try: + ran = check() + if check in (check_table_cells, check_glue_cells) and expected_cells >= 0: + assert ran == expected_cells, ( + f"ran {ran} cells, the entry declares {expected_cells}" + ) + except BaseException: + import traceback + + print(f"[rank {RANK}] FAILED {check.__name__}", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + # Exit hard: a rank that leaves an exchange early leaves its peers polling. + os._exit(13) + dist.barrier(group=HOST_GROUP) + torch.cuda.synchronize() + print(f"[rank {RANK}] {entry}: {len(checks)} checks passed", flush=True) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks(entry: str, expected_cells: int) -> None: + """One worker per device in CUDA_VISIBLE_DEVICES; every session is killed on a failure or the deadline.""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the devices this run owns" + devices = [d for d in visible.split(",") if d.strip()] + world_size = len(devices) + assert world_size >= 2, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); an exchange needs 2" + ) + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + ) + # The children rendezvous through torch.distributed's env://, not the launching step's PMI namespace. + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + print(f"[launcher] spawning {world_size} ranks for {entry}, rendezvous on :{port}", flush=True) + children = [ + subprocess.Popen( + [ + sys.executable, + str(Path(__file__).resolve()), + entry, + str(expected_cells), + _WORKER_FLAG, + ], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + for r in range(world_size) + ] + + def _kill_all() -> None: + for p in children: + try: + os.killpg(p.pid, signal.SIGKILL) + except ProcessLookupError: + pass + for p in children: + p.wait() + + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" + break + for r in sorted(live): + code = children[r].poll() + if code is None: + continue + live.discard(r) + if code != 0: + failure = f"rank {r} exited {code}" + break + time.sleep(0.5) + if failure is not None: + _kill_all() + raise AssertionError(f"the {world_size}-rank {entry} run failed: {failure}") + + +if __name__ == "__main__": + args = [a for a in sys.argv[1:] if a != _WORKER_FLAG] + entry_name = args[0] + cells = int(args[1]) if len(args) > 1 else -1 + if _WORKER_FLAG in sys.argv: + sys.exit(_run_one_rank(entry_name, cells)) + _spawn_ranks(entry_name, cells) + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_rank_job.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_rank_job.py new file mode 100644 index 000000000000..8adeed47174f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_rank_job.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Run `_draft_step_op_matrix.py` for one entry from pytest, and turn its exit code into an assertion. + +The matrix owns the hard parts -- one rank per device, its own process group, a deadline it enforces by killing +every rank's session -- so this only selects the devices and starts it in a fresh interpreter, which must carry +no CUDA or torch.distributed state of its own (a pytest process that imported tensorrt_llm does). Started by +file path: the ranks reach the target by absolute import, and this tree is not a package. +""" + +from __future__ import annotations + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path + +WORLD_SIZE = 4 +"""Ranks of the single-node run: the 16-rank geometry needs four nodes and an external launcher.""" + +_LAUNCHER_GRACE_S = 300 + + +def _devices() -> str: + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible: + devices = [d for d in visible.split(",") if d.strip()] + else: + import torch + + devices = [str(i) for i in range(torch.cuda.device_count())] + assert len(devices) >= WORLD_SIZE, ( + f"the draft-step exchange runs over {WORLD_SIZE} ranks; " + f"only {len(devices)} device(s) are visible" + ) + return ",".join(devices[:WORLD_SIZE]) + + +def run(entry: str, expected_cells: int) -> None: + """Run the matrix for `entry` over WORLD_SIZE devices, expecting it to drive `expected_cells` cells.""" + launcher = Path(__file__).resolve().with_name("_draft_step_op_matrix.py") + env = dict(os.environ, CUDA_VISIBLE_DEVICES=_devices()) + # The ranks import tensorrt_llm absolutely: the package this process imports, installed or not. + package = importlib.util.find_spec("tensorrt_llm") + assert package is not None and package.origin is not None, "tensorrt_llm is not importable" + package_parent = str(Path(package.origin).resolve().parents[1]) + env["PYTHONPATH"] = os.pathsep.join(p for p in (package_parent, env.get("PYTHONPATH", "")) if p) + spec = importlib.util.spec_from_file_location(launcher.stem, launcher) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + completed = subprocess.run( + [sys.executable, str(launcher), entry, str(expected_cells)], + env=env, + capture_output=True, + text=True, + timeout=module.DEADLINE_S + _LAUNCHER_GRACE_S, + ) + if completed.returncode != 0: + raise AssertionError( + f"{launcher.name} {entry} exited {completed.returncode}\n" + f"--- stdout ---\n{completed.stdout}\n" + f"--- stderr ---\n{completed.stderr}" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_ehproj_table.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_ehproj_table.py new file mode 100644 index 000000000000..b536ec7d560e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_ehproj_table.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the glm_draft_ehproj_table entry's multi-rank certification. + +The cells, a draft loop of the target's shape (eager and as CUDA-graph replays) and the guards run inside +`_draft_step_op_matrix.py` beside this file, one process per rank over a 4-rank group; see that file for the +16-rank run over four nodes. +""" + +import _draft_step_rank_job +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec.draft_ehproj_table import ( # noqa: E501 + glm_draft_ehproj_table as op, +) +from tensorrt_llm._utils import get_sm_version + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" +) + + +# Each run spawns its own 4 rank processes over the visible devices: under xdist several workers would fight +# for the same GPUs. +@pytest.mark.no_xdist +def test_draft_ehproj_table_op_matrix() -> None: + assert op.CELLS, "the entry declares no cells" + _draft_step_rank_job.run("draft_ehproj_table", expected_cells=len(op.CELLS)) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py new file mode 100644 index 000000000000..3e211989b671 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm_draft_glue_rmsnorm_cat entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec.draft_glue_rmsnorm_cat import ( # noqa: E501 + glm_draft_glue_rmsnorm_cat as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402,F401 — needs the path declared above + assert_within_ulp, + certifying, + validating, +) + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" +) + +HIDDEN = 6144 +EPS = 1e-5 + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """One RMSNorm's worth of rounding differences. + + Kernel and reference normalize the same BF16 rows and round once to BF16. They differ in the order of the 6144 + squares in the FP32 sum and in the reciprocal square root (the kernel's approximate `rsqrt`, a few FP32 ulp, the + reference's correctly rounded one): together a few 1e-7 relative on the scale factor, which moves the BF16 + rounding of an output only when it lies that close to a rounding midpoint, by one ulp of that element. + + One BF16 ulp of an element in the row maximum's binade is up to twice the band's unit (2**-8 of the row + maximum). So: 2 ulp element-wise against the row scale, 0.25 ulp of relative RMS. A wrong row, column or weight + is off by the size of the values themselves. + """ + assert_within_ulp(out, ref, element_ulp=2.0, rms_ulp=0.25) + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + rows, vocab = spec["rows"], spec["vocab"] + table = torch.randn(vocab, HIDDEN, generator=g, device="cuda", dtype=torch.bfloat16) + hidden = 3.0 * torch.randn(rows, HIDDEN, generator=g, device="cuda") + if spec["ids"] == "valid": + ids = torch.randint(0, vocab, (rows,), generator=g, device="cuda", dtype=torch.int32) + else: + # int64 ids: -1 padding, the first id past the table, valid ids, a large negative and a far one. + candidates = [-1, vocab, 5, -7, vocab + 100, 0, vocab - 1, -1] + ids = torch.tensor(candidates[:rows], dtype=torch.int64, device="cuda") + enorm = (1.0 + 0.1 * torch.randn(HIDDEN, generator=g, device="cuda")).to(torch.bfloat16) + hnorm = (1.0 + 0.1 * torch.randn(HIDDEN, generator=g, device="cuda")).to(torch.bfloat16) + lo, hi = spec["window"] + return table, ids, hidden.to(torch.bfloat16), enorm, hnorm, EPS, lo, hi + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + out = op(*args, True) + assert out.is_contiguous() and out.shape == (cell.spec["rows"], args[7] - args[6]) + _compare(out, op.reference(*args, True)) + + +def test_window_is_flashinfer_rmsnorm_bitwise() -> None: + """Where FlashInfer runs its CuTe RMSNorm, the window is that kernel's output on each half, bit for bit.""" + import flashinfer.norm + + device = torch.device("cuda") + if flashinfer.norm._use_cuda_norm(device): + pytest.skip("flashinfer.norm.rmsnorm runs its CUDA JIT kernel on this device") + for rank in (0, 15): + table, ids, hidden, enorm, hnorm, eps, lo, hi = _build( + dict(rows=8, vocab=4096, window=(768 * rank, 768 * (rank + 1)), ids="valid"), seed=rank + ) + embeds = torch.nn.functional.embedding(ids.long(), table) + enormed = flashinfer.norm.rmsnorm(embeds, enorm, eps) + hnormed = flashinfer.norm.rmsnorm(hidden, hnorm, eps) + unfused = torch.cat([enormed, hnormed], dim=-1) + out = op(table, ids, hidden, enorm, hnorm, eps, lo, hi, True) + torch.testing.assert_close(out, unfused[:, lo:hi], rtol=0, atol=0) + + +def test_guard_refuses_a_window_inside_a_vector() -> None: + """The store predicate tests each 8-column vector's first column only.""" + table, ids, hidden, enorm, hnorm, eps, _, _ = _build(op.CELLS[0].spec, seed=1) + with pytest.raises(AssertionError, match="8-column aligned"): + with validating(op): + op(table, ids, hidden, enorm, hnorm, eps, 4, 772, True) + + +def test_guard_refuses_a_misaligned_hidden_row() -> None: + """The kernel is compiled for 16-byte aligned rows.""" + table, ids, hidden, enorm, hnorm, eps, lo, hi = _build(op.CELLS[2].spec, seed=2) + storage = torch.zeros(8, HIDDEN + 8, dtype=torch.bfloat16, device="cuda") + shifted = storage[:, 1 : HIDDEN + 1] + shifted.copy_(hidden) + with pytest.raises(AssertionError, match="16-byte aligned"): + with validating(op): + op(table, ids, shifted, enorm, hnorm, eps, lo, hi, True) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_sample_glue.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_sample_glue.py new file mode 100644 index 000000000000..faf00183c753 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_sample_glue.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the glm_draft_sample_glue entry's multi-rank certification. + +The cells, a draft loop of the target's shape (eager and as CUDA-graph replays) and the guards run inside +`_draft_step_op_matrix.py` beside this file, one process per rank over a 4-rank group; see that file for the +16-rank run over four nodes. +""" + +import _draft_step_rank_job +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec.draft_sample_glue import ( # noqa: E501 + glm_draft_sample_glue as op, +) +from tensorrt_llm._utils import get_sm_version + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" +) + + +# Each run spawns its own 4 rank processes over the visible devices: under xdist several workers would fight +# for the same GPUs. +@pytest.mark.no_xdist +def test_draft_sample_glue_op_matrix() -> None: + assert op.CELLS, "the entry declares no cells" + _draft_step_rank_job.run("draft_sample_glue", expected_cells=len(op.CELLS)) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_token_table.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_token_table.py new file mode 100644 index 000000000000..54b0750306bf --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_token_table.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the construction-time operands of the glm_draft_ehproj_table entry.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec._draft_token_table import ( # noqa: E501 + build_draft_token_table, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" +) + +HIDDEN = 6144 +EPS = 1e-5 + + +@pytest.mark.parametrize("tp_size,tp_rank", [(16, 0), (16, 11), (4, 3)]) +def test_table_is_the_token_half_of_eh_proj(tp_size, tp_rank) -> None: + """Each table row is eh_proj's embedding half of one token of the rank's shard. + + Against an FP32 RMSNorm and product, rounded where the builder rounds (the normalized row and the product to + BF16), the rows differ only where the kernel's approximate reciprocal square root moves the normalized row's + BF16 rounding (one ulp of an input, carried through the product) and where the order of a 6144-term FP32 sum + moves the product's (one ulp of the element): 2 ulp per element, 0.25 ulp of relative RMS. The hidden-half + rows are a slice of the weight, compared exactly. + """ + g = torch.Generator(device="cuda").manual_seed(tp_size * 100 + tp_rank) + vocab_shard = 96 + weight = (torch.randn(HIDDEN, 2 * HIDDEN, generator=g, device="cuda") * HIDDEN**-0.5).to( + torch.bfloat16 + ) + embed = torch.randn(vocab_shard * tp_size, HIDDEN, generator=g, device="cuda").to( + torch.bfloat16 + ) + enorm = (1.0 + 0.1 * torch.randn(HIDDEN, generator=g, device="cuda")).to(torch.bfloat16) + + table, hidden_rows = build_draft_token_table( + weight, embed, enorm, EPS, tp_rank, tp_size, vocab_shard + ) + + rows = embed[tp_rank * vocab_shard : (tp_rank + 1) * vocab_shard].float() + rstd = torch.rsqrt(rows.pow(2).mean(dim=-1, keepdim=True) + EPS) + normed = ((rows * rstd) * enorm.float()).to(torch.bfloat16).float() + with true_fp32_matmul(): + expected = (normed @ weight[:, :HIDDEN].float().t()).to(torch.bfloat16) + assert table.shape == (vocab_shard, HIDDEN) and table.dtype == torch.bfloat16 + assert_within_ulp(table, expected, element_ulp=2.0, rms_ulp=0.25) + share = HIDDEN // tp_size + torch.testing.assert_close( + hidden_rows, weight[tp_rank * share : (tp_rank + 1) * share, HIDDEN:], rtol=0, atol=0 + ) + assert hidden_rows.is_contiguous() diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_gather_last_accepted_rows.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_gather_last_accepted_rows.py new file mode 100644 index 000000000000..3edc7f4423e5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_gather_last_accepted_rows.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm_gather_last_accepted_rows entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec.gather_last_accepted_rows import ( # noqa: E501 + glm_gather_last_accepted_rows as op, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — needs the path declared above + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" +) + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + requests, per_request, cols = spec["requests"], spec["rows_per_request"], spec["cols"] + rows = torch.randn(requests * per_request, cols, generator=g, device="cuda").to(torch.bfloat16) + second = None + if spec["second"]: + second = torch.randn(requests * per_request, cols, generator=g, device="cuda").to( + torch.bfloat16 + ) + order = spec.get("batch_order", tuple(range(requests))) + batch_indices = torch.tensor(order, dtype=torch.int32, device="cuda") + num_accepted = torch.tensor(spec["accepted"], dtype=torch.int32, device="cuda") + return rows, batch_indices, num_accepted, per_request, 0, second + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + """A gather moves bytes: the op and the reference agree bit for bit.""" + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + out = op(*args) + ref = op.reference(*args) + assert len(out) == len(ref) == (3 if cell.spec["second"] else 2) + for got, expected in zip(out, ref): + torch.testing.assert_close(got, expected, rtol=0, atol=0) + + +def test_row_offset_shifts_the_selection() -> None: + """Rows after `row_offset` context tokens: the index moves by the offset, the gathers with it.""" + rows, batch_indices, num_accepted, per_request, _, second = _build(op.CELLS[0].spec, seed=3) + padded = torch.cat([torch.zeros(5, rows.shape[1], dtype=rows.dtype, device="cuda"), rows]) + padded_second = torch.cat([torch.zeros_like(padded[:5]), second]) + with validating(op): + out = op(padded, batch_indices, num_accepted, per_request, 5, padded_second) + ref = op.reference(rows, batch_indices, num_accepted, per_request, 0, second) + torch.testing.assert_close(out[0], ref[0] + 5, rtol=0, atol=0) + torch.testing.assert_close(out[1], ref[1], rtol=0, atol=0) + torch.testing.assert_close(out[2], ref[2], rtol=0, atol=0) + + +def test_guard_refuses_a_count_of_zero() -> None: + """A request with no accepted row would select the row before its own.""" + rows, batch_indices, _, per_request, offset, second = _build(op.CELLS[0].spec, seed=4) + zero = torch.zeros(1, dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match=r"\[1, rows_per_request\]"): + with validating(op): + op(rows, batch_indices, zero, per_request, offset, second) + + +def test_guard_refuses_a_wider_second() -> None: + """The kernel copies the columns of `rows`: a wider `second` keeps garbage in its tail.""" + rows, batch_indices, num_accepted, per_request, offset, _ = _build(op.CELLS[0].spec, seed=5) + wide = torch.zeros(rows.shape[0], rows.shape[1] + 8, dtype=rows.dtype, device="cuda") + with pytest.raises(AssertionError, match="columns of rows"): + with validating(op): + op(rows, batch_indices, num_accepted, per_request, offset, wide) + + +def test_guard_refuses_a_strided_count_vector() -> None: + """The kernel reads the counts as a dense vector.""" + rows, _, _, per_request, offset, second = _build(op.CELLS[2].spec, seed=6) + batch_indices = torch.arange(4, dtype=torch.int32, device="cuda") + strided = torch.ones(8, dtype=torch.int32, device="cuda")[::2] + with pytest.raises(AssertionError, match="contiguous vector"): + with validating(op): + op(rows, batch_indices, strided, per_request, offset, second) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_verify_head_accept.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_verify_head_accept.py new file mode 100644 index 000000000000..9cef188ae5b6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_verify_head_accept.py @@ -0,0 +1,201 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the glm_verify_head_accept entry.""" + +import functools +import zlib +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec.verify_head_accept import ( # noqa: E501 + glm_verify_head_accept as op, +) +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + SpecWorkerBase, +) +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — needs the path declared above + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" +) + +VOCAB = 154880 +STRIDES = ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, +) + + +def _spec_worker_rng_state(device): + """The pool and counter the one-model spec worker builds (fixed seed, identical on every rank).""" + holder = SimpleNamespace(_force_accept_rng_pool=None, _force_accept_rng_counter=None) + SpecWorkerBase._ensure_force_accept_rng_state(holder, device) + return holder._force_accept_rng_pool, holder._force_accept_rng_counter + + +def _build(spec, seed): + """One cell's operands. + + Target tokens, draft tokens with `matches[g]` leading tokens equal to request g's targets, the spec worker's + pool, a counter at a seed-dependent value and, with `rows`, a sentinel-filled accepted_tokens. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + num_contexts, num_gens, draft_len = spec["num_contexts"], spec["num_gens"], spec["draft_len"] + num_targets = num_contexts + num_gens * (draft_len + 1) + target = torch.randint(0, VOCAB - 1, (num_targets,), generator=g, device="cuda") + gen_target = target[num_contexts:].view(num_gens, draft_len + 1) + draft = torch.randint(0, VOCAB, (num_gens, draft_len), generator=g, device="cuda") + for request, matched in enumerate(spec["matches"]): + draft[request, :matched] = gen_target[request, :matched] + if matched < draft_len: + # Targets stay below VOCAB - 1, so this is a valid token that differs from the target. + draft[request, matched] = gen_target[request, matched] + 1 + pool, counter = _spec_worker_rng_state(torch.device("cuda")) + counter.fill_(seed % 997) + rows = None + if spec["rows"]: + rows = torch.full( + (num_contexts + num_gens, draft_len + 1), -7, dtype=torch.int32, device="cuda" + ) + return ( + target.to(spec["target_dtype"]).contiguous(), + draft.to(torch.int32).contiguous(), + rows, + pool, + counter, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + """Integer outputs: the op and the ATen chain agree exactly.""" + spec = cell.spec + target, draft, rows, pool, counter = _build(spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + mode, base, maximum, frac = spec["force"] + rows_before = rows.clone() if rows is not None else None + counter_before = counter.clone() + certifying(op) + contexts = spec["num_contexts"] + counts = op(target, draft, contexts, rows, pool, counter, mode, base, maximum, frac, *STRIDES) + ref_counts, ref_rows, ref_counter = op.reference( + target, + draft, + spec["num_contexts"], + rows_before, + pool, + counter_before, + mode, + base, + maximum, + frac, + *STRIDES, + ) + torch.testing.assert_close(counts, ref_counts, rtol=0, atol=0) + if rows is not None: + torch.testing.assert_close(rows, ref_rows, rtol=0, atol=0) + if ref_counter is not None: + torch.testing.assert_close(counter, ref_counter, rtol=0, atol=0) + else: + torch.testing.assert_close(counter, counter_before, rtol=0, atol=0) + + +@pytest.mark.parametrize("forced", [2.91, 2.5, 6.1, 3.0]) +def test_forced_count_matches_the_spec_worker(forced) -> None: + """The op's forced counts equal the spec worker's over consecutive steps. + + `SpecWorkerBase._apply_force_accepted_tokens` runs on its own counter, the op on a copy of it with the same + pool and the arguments the target derives from the forced value as that method does. + """ + draft_len, num_gens = 7, 4 + worker = SimpleNamespace( + force_num_accepted_tokens=forced, + _force_accept_rng_pool=None, + _force_accept_rng_counter=None, + ) + worker._ensure_force_accept_rng_state = functools.partial( + SpecWorkerBase._ensure_force_accept_rng_state, worker + ) + worker._ensure_force_accept_rng_state(torch.device("cuda")) + op_counter = worker._force_accept_rng_counter.clone() + int_part = int(forced) + frac_part = forced - int_part + max_total = draft_len + 1 + base_total = min(int_part + 1, max_total) + mode = 2 if frac_part > 0.0 and base_total < max_total else 1 + target, draft, _, _, _ = _build( + dict( + num_contexts=0, + num_gens=num_gens, + draft_len=draft_len, + target_dtype=torch.int32, + matches=(0,) * num_gens, + rows=False, + ), + seed=17, + ) + for _ in range(64): + strict = torch.ones(num_gens, dtype=torch.int32, device="cuda") + expected = SpecWorkerBase._apply_force_accepted_tokens(worker, strict, 0, draft_len) + counts = op( + target, + draft, + 0, + None, + worker._force_accept_rng_pool, + op_counter, + mode, + base_total, + max_total, + frac_part if mode == 2 else 0.0, + *STRIDES, + ) + torch.testing.assert_close(counts, expected, rtol=0, atol=0) + torch.testing.assert_close(op_counter, worker._force_accept_rng_counter, rtol=0, atol=0) + + +def test_cuda_graph_replays_advance_the_counter() -> None: + """Inside the decode graph every replay draws the next pool entries, as the eager calls would.""" + spec = op.CELLS[0].spec + target, draft, _, pool, counter = _build(spec, seed=5) + force = spec["force"] + # One eager call first, as the engine's warm-up does before capture. + op(target, draft, 0, None, pool, counter, *force, *STRIDES) + expected_counter = counter.clone() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_counts = op(target, draft, 0, None, pool, counter, *force, *STRIDES) + for _ in range(8): + expected, _, expected_counter = op.reference( + target, draft, 0, None, pool, expected_counter, *force, *STRIDES + ) + graph.replay() + torch.testing.assert_close(graph_counts, expected, rtol=0, atol=0) + torch.testing.assert_close(counter, expected_counter, rtol=0, atol=0) + + +def test_guard_refuses_a_wider_rows_buffer() -> None: + """Rows of the maximum draft length on a step that verifies fewer. + + The op's size check accepts them, and the rows are written with the wrong row stride. + """ + target, draft, _, pool, counter = _build(op.CELLS[0].spec, seed=3) + wide = torch.zeros(1, 12, dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match="exactly"): + with validating(op): + op(target, draft, 0, wide, pool, counter, 0, 0, 0, 0.0, *STRIDES) + + +def test_guard_refuses_a_forced_count_past_the_draft_window() -> None: + target, draft, _, pool, counter = _build(op.CELLS[0].spec, seed=4) + with pytest.raises(AssertionError, match="force_max == draft_len"): + with validating(op): + op(target, draft, 0, None, pool, counter, 1, 9, 9, 0.0, *STRIDES) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py index 81256a1c2c97..42d1d187f016 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -338,3 +338,74 @@ def test_qwen4_exp_deepstack_vision_tower_does_not_match(): def test_qwen4_exp_on_two_ranks_does_not_match(): config = _model_config(_qwen4_exp_config(), world_size=2, tp_size=2) assert modeling_v2_resolve(config, _llm_args("auto")) is None + + +# --- glm_moe_dsa: tp16 on sm107, MTP-Eagle over a linear chain at any length ---- + +_GLM52_TARGET = ( + "models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.modeling.ModelingV2Glm52Nvfp4Sm107Tp16" +) +_TP16 = dict(world_size=16, tp_size=16, moe_ep_size=1, moe_tp_size=16) + + +@pytest.fixture +def _on_sm107(monkeypatch): + """Route as if this were a Rubin GPU.""" + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (10, 7)) + + +def _glm52_config(**overrides): + """The shape GLM-5.2-NVFP4's own config.json declares.""" + fields = dict( + architectures=["GlmMoeDsaForCausalLM"], + model_type="glm_moe_dsa", + num_hidden_layers=78, + hidden_size=6144, + n_routed_experts=256, + q_lora_rank=2048, + ) + fields.update(overrides) + return PretrainedConfig(**fields) + + +@pytest.mark.parametrize("mode", ["auto", "require"]) +@pytest.mark.parametrize( + "spec_config", + [None, MTPDecodingConfig(max_draft_len=7), MTPDecodingConfig(max_draft_len=3)], + ids=["no-spec", "mtp7", "mtp3"], +) +def test_glm52_tp16_matches(_on_sm107, spec_config, mode): + """Any draft length: the draft loop replays the one MTP module per step.""" + config = _model_config(_glm52_config(), spec_config=spec_config, **_TP16) + assert _qualified(modeling_v2_resolve(config, _llm_args(mode))) == _GLM52_TARGET + + +@pytest.mark.parametrize( + "config_kwargs, mapping_kwargs, spec_config, missed", + [ + (dict(num_hidden_layers=92), _TP16, None, r"shape \(92, 6144, 256, 2048\)"), + (dict(), dict(_TP16, enable_attention_dp=True), None, "attention_dp=True"), + (dict(), dict(_TP16, moe_ep_size=16, moe_tp_size=1), None, "moe_ep=16"), + (dict(), dict(world_size=8, tp_size=8, moe_tp_size=8), None, "world_size=8"), + (dict(), _TP16, MTPDecodingConfig(max_draft_len=3, use_mtp_vanilla=True), "MTP"), + (dict(), _TP16, _NGRAM, "NGRAM"), + ], + ids=["other-shape", "attention-dp", "expert-parallel", "tp8", "vanilla-mtp", "ngram"], +) +def test_glm52_near_misses_do_not_match( + _on_sm107, config_kwargs, mapping_kwargs, spec_config, missed +): + config = _model_config( + _glm52_config(**config_kwargs), spec_config=spec_config, **mapping_kwargs + ) + assert modeling_v2_resolve(config, _llm_args("auto")) is None + with pytest.raises(ValueError, match=missed): + modeling_v2_resolve(config, _llm_args("require")) + + +def test_glm52_on_sm103_does_not_match(): + """The target's kernels are built for sm107; a GB300 falls back.""" + config = _model_config(_glm52_config(), **_TP16) + assert modeling_v2_resolve(config, _llm_args("auto")) is None + with pytest.raises(ValueError, match=r"compute capability \(10, 3\)"): + modeling_v2_resolve(config, _llm_args("require"))